diff --git a/cmd/armadactl/cmd/cancel.go b/cmd/armadactl/cmd/cancel.go index 5d33ee8d3d6..1a75906fae2 100644 --- a/cmd/armadactl/cmd/cancel.go +++ b/cmd/armadactl/cmd/cancel.go @@ -20,9 +20,9 @@ func cancelCmd() *cobra.Command { cmd.AddCommand( cancelJobCmd(), cancelJobSetCmd(), - cancelExecutorCmd(), - cancelNodeCmd(), - cancelQueueCmd(), + cancelExecutorCmd(armadactl.New()), + cancelNodeCmd(armadactl.New()), + cancelQueueCmd(armadactl.New()), ) return cmd } @@ -66,23 +66,13 @@ func cancelJobSetCmd() *cobra.Command { return cmd } -func cancelExecutorCmd() *cobra.Command { - a := armadactl.New() +func cancelExecutorCmd(a *armadactl.App) *cobra.Command { cmd := &cobra.Command{ Use: "executor ", Short: "Cancels jobs on executor.", Long: `Cancels jobs on executor with provided executor name, priority classes, queues, and pools.`, Args: cobra.ExactArgs(1), PreRunE: func(cmd *cobra.Command, args []string) error { - all, err := cmd.Flags().GetBool("all-priority-classes") - if err != nil { - return fmt.Errorf("error reading all-priority-classes flag: %s", err) - } - if !all { - if err := cmd.MarkFlagRequired("priority-classes"); err != nil { - return fmt.Errorf("error marking priority-class flag as required: %s", err) - } - } return initParams(cmd, a.Params) }, RunE: func(cmd *cobra.Command, args []string) error { @@ -92,10 +82,6 @@ func cancelExecutorCmd() *cobra.Command { if err != nil { return fmt.Errorf("error reading priority-class selection: %s", err) } - allPriorityClasses, _ := cmd.Flags().GetBool("all-priority-classes") - if allPriorityClasses { - priorityClasses = nil - } queues, err := cmd.Flags().GetStringSlice("queues") if err != nil { @@ -121,13 +107,7 @@ func cancelExecutorCmd() *cobra.Command { "priority-classes", "p", []string{}, - "Cancel jobs on executor matching the specified priority classes. Provided priority classes should be comma separated, as in the following example: armada-default,armada-preemptible.", - ) - cmd.Flags().BoolP( - "all-priority-classes", - "a", - false, - "Cancel jobs on executor for all priority classes.", + "Cancel jobs on executor matching the specified priority classes, comma separated (e.g. armada-default,armada-preemptible). If no priority classes are provided, jobs across all priority classes will be cancelled.", ) cmd.Flags().StringSlice( "pools", @@ -137,23 +117,13 @@ func cancelExecutorCmd() *cobra.Command { return cmd } -func cancelNodeCmd() *cobra.Command { - a := armadactl.New() +func cancelNodeCmd(a *armadactl.App) *cobra.Command { cmd := &cobra.Command{ Use: "node ", Short: "Cancels jobs on node for specified executor.", Long: `Cancels jobs on node for executor with provided executor name, priority classes and queues.`, Args: cobra.ExactArgs(1), PreRunE: func(cmd *cobra.Command, args []string) error { - all, err := cmd.Flags().GetBool("all-priority-classes") - if err != nil { - return fmt.Errorf("error reading all-priority-classes flag: %s", err) - } - if !all { - if err := cmd.MarkFlagRequired("priority-classes"); err != nil { - return fmt.Errorf("error marking priority-class flag as required: %s", err) - } - } if err := cmd.MarkFlagRequired("executor"); err != nil { return fmt.Errorf("error marking executor flag as required: %s", err) } @@ -166,10 +136,6 @@ func cancelNodeCmd() *cobra.Command { if err != nil { return fmt.Errorf("error reading priority-class selection: %s", err) } - allPriorityClasses, _ := cmd.Flags().GetBool("all-priority-classes") - if allPriorityClasses { - priorityClasses = nil - } queues, err := cmd.Flags().GetStringSlice("queues") if err != nil { @@ -195,13 +161,7 @@ func cancelNodeCmd() *cobra.Command { "priority-classes", "p", []string{}, - "Cancel jobs on node for specified executor matching the specified priority classes. Provided priority classes should be comma separated, as in the following example: armada-default,armada-preemptible.", - ) - cmd.Flags().BoolP( - "all-priority-classes", - "a", - false, - "Preempt jobs on executor for all priority classes.", + "Cancel jobs on node for specified executor matching the specified priority classes, comma separated (e.g. armada-default,armada-preemptible). If no priority classes are provided, jobs across all priority classes will be cancelled.", ) cmd.Flags().StringP( "executor", @@ -212,8 +172,7 @@ func cancelNodeCmd() *cobra.Command { return cmd } -func cancelQueueCmd() *cobra.Command { - a := armadactl.New() +func cancelQueueCmd(a *armadactl.App) *cobra.Command { cmd := &cobra.Command{ Use: "queues ...", Short: "Cancels jobs on queues.", @@ -223,15 +182,6 @@ func cancelQueueCmd() *cobra.Command { if err := cmd.MarkFlagRequired("job-states"); err != nil { return err } - all, err := cmd.Flags().GetBool("all-priority-classes") - if err != nil { - return fmt.Errorf("error reading all-priority-classes flag: %s", err) - } - if !all { - if err := cmd.MarkFlagRequired("priority-classes"); err != nil { - return fmt.Errorf("error marking priority-class flag as required: %s", err) - } - } return initParams(cmd, a.Params) }, RunE: func(cmd *cobra.Command, queues []string) error { @@ -273,10 +223,6 @@ func cancelQueueCmd() *cobra.Command { if err != nil { return fmt.Errorf("error reading priority-classes flag: %s", err) } - allPriorityClasses, _ := cmd.Flags().GetBool("all-priority-classes") - if allPriorityClasses { - priorityClasses = nil - } pools, err := cmd.Flags().GetStringSlice("pools") if err != nil { @@ -313,13 +259,7 @@ func cancelQueueCmd() *cobra.Command { "priority-classes", "p", []string{}, - "Jobs matching the provided priority classes will be cancelled.", - ) - cmd.Flags().BoolP( - "all-priority-classes", - "a", - false, - "Preempt jobs on executor for all priority classes.", + "Jobs matching the provided priority classes will be cancelled. If no priority classes are provided, jobs across all priority classes will be cancelled.", ) cmd.Flags().StringSlice( "pools", diff --git a/cmd/armadactl/cmd/cancel_test.go b/cmd/armadactl/cmd/cancel_test.go index 609ec4ce6ff..a2cb9cc9f00 100644 --- a/cmd/armadactl/cmd/cancel_test.go +++ b/cmd/armadactl/cmd/cancel_test.go @@ -1,277 +1,280 @@ package cmd import ( - "io" "testing" - "github.com/spf13/cobra" "github.com/stretchr/testify/require" "github.com/armadaproject/armada/internal/armadactl" + "github.com/armadaproject/armada/pkg/api" ) -func TestCancel(t *testing.T) { +func TestCancelExecutor(t *testing.T) { tests := map[string]struct { - Flags []flag - jobId string - queue string - jobSet string + flags []flag + want executorCall }{ - "default flags": {nil, "", "", ""}, - "valid job-id": {[]flag{{"job-id", "jobId1"}}, "jobId1", "", ""}, - "valid queue": {[]flag{{"queue", "queue1,jobSet1"}}, "", "queue1", "jobSet1"}, - "valid job-set": {[]flag{{"job-set", "jobSet1"}}, "", "", "jobSet1"}, - } - for name, test := range tests { - t.Run(name, func(t *testing.T) { - a := armadactl.New() - cmd := cancelCmd() - - cmd.PreRunE = func(cmd *cobra.Command, args []string) error { - a.Out = io.Discard - - if len(test.jobId) > 0 { - jobIdFlag, err1 := cmd.Flags().GetString("job-id") - require.Error(t, err1) - require.Equal(t, test.jobId, jobIdFlag) - } - if len(test.queue) > 0 { - queueFlag, err1 := cmd.Flags().GetString("queue") - jobSetFlag, err2 := cmd.Flags().GetString("job-set") - require.Error(t, err1) - require.Error(t, err2) - require.Equal(t, test.queue, queueFlag) - require.Equal(t, test.jobSet, jobSetFlag) - } - if len(test.jobSet) > 0 { - jobSetFlag, err1 := cmd.Flags().GetString("job-set") - require.Error(t, err1) - require.Equal(t, test.jobSet, jobSetFlag) - } - return nil - } - }) - } -} - -func TestCancelQueue(t *testing.T) { - tests := map[string]struct { - Flags []flag - jobStates []string - selectors []string - priorityClasses []string - inverse bool - onlyCordoned bool - dryRun bool - }{ - "default flags": {nil, []string{}, []string{}, []string{}, false, false, false}, - "valid selectors": {[]flag{{"selectors", "armadaproject.io/priority=high,armadaproject.io/category=critical"}}, []string{}, []string{"armadaproject.io/priority=high", "armadaproject.io/category=critical"}, []string{}, false, false, false}, - "valid job-states 1": {[]flag{{"job-states", "queued"}}, []string{"queued"}, []string{}, []string{}, false, false, false}, - "valid job-states 2": {[]flag{{"job-states", "queued,leased,pending,running"}}, []string{"queued", "leased", "pending", "running"}, []string{}, []string{}, false, false, false}, - "valid priority-classes 1": {[]flag{{"priority-classes", "armada-default"}}, []string{}, []string{}, []string{"armada-default"}, false, false, false}, - "valid priority-classes 2": {[]flag{{"priority-classes", "armada-default,armada-preemptible"}}, []string{}, []string{}, []string{"armada-default", "armada-preemptible"}, false, false, false}, - "valid multiple flags": { - []flag{{"selectors", "armadaproject.io/priority=high,armadaproject.io/category=critical"}, {"job-states", "queued,leased,pending,running"}, {"priority-classes", "armada-default,armada-preemptible"}}, - []string{"queued", "leased", "pending", "running"}, - []string{"armadaproject.io/priority=high", "armadaproject.io/category=critical"}, - []string{"armada-default", "armada-preemptible"}, - true, true, true, + // Omitting priority-classes means all priority classes, which the + // executor API represents as an empty slice. An unnarrowed queue + // selection expands to every queue. + "without priority-classes": { + flags: nil, + want: executorCall{ + executor: "test-executor", + queues: []string{"queue-a", "queue-b"}, + priorityClasses: []string{}, + pools: []string{}, + }, }, - } - for name, test := range tests { - t.Run(name, func(t *testing.T) { - a := armadactl.New() - cmd := cancelQueueCmd() - - cmd.PreRunE = func(cmd *cobra.Command, args []string) error { - a.Out = io.Discard - - if len(test.jobStates) > 0 { - jobStatesFlag, err := cmd.Flags().GetString("job-states") - require.NoError(t, err) - require.Equal(t, test.jobStates, jobStatesFlag) - } - if len(test.selectors) > 0 { - selectorsFlag, err := cmd.Flags().GetString("selectors") - require.Error(t, err) - require.Equal(t, test.selectors, selectorsFlag) - } - if len(test.priorityClasses) > 0 { - priorityClassesFlag, err := cmd.Flags().GetString("priority-classes") - require.Error(t, err) - require.Equal(t, test.priorityClasses, priorityClassesFlag) - } - - inverseValue, err := cmd.Flags().GetBool("inverse") - require.NoError(t, err) - require.Equal(t, test, inverseValue) - - onlyCordonedValue, err := cmd.Flags().GetBool("only-cordoned") - require.NoError(t, err) - require.Equal(t, test, onlyCordonedValue) - - dryRunValue, err := cmd.Flags().GetBool("dry-run") - require.NoError(t, err) - require.Equal(t, test, dryRunValue) - - return nil - } - }) - } -} - -func TestCancelExecutorAllPriorityClasses(t *testing.T) { - tests := map[string]struct { - flags []flag - expectError bool - }{ - "with all-priority-classes flag set": { - flags: []flag{{"all-priority-classes", "true"}}, - expectError: false, + "with a single priority class": { + flags: []flag{{"priority-classes", "armada-default"}}, + want: executorCall{ + executor: "test-executor", + queues: []string{"queue-a", "queue-b"}, + priorityClasses: []string{"armada-default"}, + pools: []string{}, + }, }, - "without all-priority-classes and without priority-classes": { - flags: nil, - expectError: true, + "with multiple priority classes": { + flags: []flag{{"priority-classes", "armada-default,armada-preemptible"}}, + want: executorCall{ + executor: "test-executor", + queues: []string{"queue-a", "queue-b"}, + priorityClasses: []string{"armada-default", "armada-preemptible"}, + pools: []string{}, + }, }, - "without all-priority-classes but with priority-classes": { - flags: []flag{{"priority-classes", "armada-default"}}, - expectError: false, + "with queues and pools": { + flags: []flag{{"queues", "queue-a"}, {"pools", "pool-1,pool-2"}}, + want: executorCall{ + executor: "test-executor", + queues: []string{"queue-a"}, + priorityClasses: []string{}, + pools: []string{"pool-1", "pool-2"}, + }, }, } for name, tc := range tests { t.Run(name, func(t *testing.T) { - cmd := cancelExecutorCmd() - cmd.PreRunE = func(cmd *cobra.Command, args []string) error { - all, err := cmd.Flags().GetBool("all-priority-classes") - if err != nil { - return err - } - if !all { - if err := cmd.MarkFlagRequired("priority-classes"); err != nil { - return err - } + a := armadactl.New() + cmd := cancelExecutorCmd(a) + + var got []executorCall + withFakeAPIs(t, a, cmd, func() { + a.Params.QueueAPI.GetAll = func() ([]*api.Queue, error) { return testQueues(), nil } + a.Params.ExecutorAPI.CancelOnExecutor = func(executor string, queues, priorityClasses, pools []string) error { + got = append(got, executorCall{executor, queues, priorityClasses, pools}) + return nil } - return nil - } - cmd.RunE = func(cmd *cobra.Command, args []string) error { - return nil - } + }) + cmd.SetArgs([]string{"test-executor"}) for _, f := range tc.flags { require.NoError(t, cmd.Flags().Set(f.name, f.value)) } - err := cmd.Execute() - if tc.expectError { - require.Error(t, err) - } else { - require.NoError(t, err) - } + + require.NoError(t, cmd.Execute()) + require.Equal(t, []executorCall{tc.want}, got) }) } } -func TestCancelNodeAllPriorityClasses(t *testing.T) { +func TestCancelNode(t *testing.T) { tests := map[string]struct { - flags []flag - expectError bool + flags []flag + want nodeCall }{ - "with all-priority-classes flag set": { - flags: []flag{{"all-priority-classes", "true"}, {"executor", "test-exec"}}, - expectError: false, + // Omitting priority-classes means all priority classes, which the node + // API represents as an empty slice. + "without priority-classes": { + flags: []flag{{"executor", "test-executor"}}, + want: nodeCall{ + node: "test-node", + executor: "test-executor", + queues: []string{"queue-a", "queue-b"}, + priorityClasses: []string{}, + }, }, - "without all-priority-classes and without priority-classes": { - flags: []flag{{"executor", "test-exec"}}, - expectError: true, + "with a single priority class": { + flags: []flag{{"executor", "test-executor"}, {"priority-classes", "armada-default"}}, + want: nodeCall{ + node: "test-node", + executor: "test-executor", + queues: []string{"queue-a", "queue-b"}, + priorityClasses: []string{"armada-default"}, + }, }, - "without all-priority-classes but with priority-classes": { - flags: []flag{{"priority-classes", "armada-default"}, {"executor", "test-exec"}}, - expectError: false, + "with multiple priority classes": { + flags: []flag{{"executor", "test-executor"}, {"priority-classes", "armada-default,armada-preemptible"}}, + want: nodeCall{ + node: "test-node", + executor: "test-executor", + queues: []string{"queue-a", "queue-b"}, + priorityClasses: []string{"armada-default", "armada-preemptible"}, + }, + }, + "with queues": { + flags: []flag{{"executor", "test-executor"}, {"queues", "queue-a"}}, + want: nodeCall{ + node: "test-node", + executor: "test-executor", + queues: []string{"queue-a"}, + priorityClasses: []string{}, + }, }, } for name, tc := range tests { t.Run(name, func(t *testing.T) { - cmd := cancelNodeCmd() - cmd.PreRunE = func(cmd *cobra.Command, args []string) error { - all, err := cmd.Flags().GetBool("all-priority-classes") - if err != nil { - return err - } - if !all { - if err := cmd.MarkFlagRequired("priority-classes"); err != nil { - return err - } - } - if err := cmd.MarkFlagRequired("executor"); err != nil { - return err + a := armadactl.New() + cmd := cancelNodeCmd(a) + + var got []nodeCall + withFakeAPIs(t, a, cmd, func() { + a.Params.QueueAPI.GetAll = func() ([]*api.Queue, error) { return testQueues(), nil } + a.Params.NodeAPI.CancelOnNode = func(node, executor string, queues, priorityClasses []string) error { + got = append(got, nodeCall{node, executor, queues, priorityClasses}) + return nil } - return nil - } - cmd.RunE = func(cmd *cobra.Command, args []string) error { - return nil - } + }) + cmd.SetArgs([]string{"test-node"}) for _, f := range tc.flags { require.NoError(t, cmd.Flags().Set(f.name, f.value)) } - err := cmd.Execute() - if tc.expectError { - require.Error(t, err) - } else { - require.NoError(t, err) - } + + require.NoError(t, cmd.Execute()) + require.Equal(t, []nodeCall{tc.want}, got) }) } } -func TestCancelQueueAllPriorityClasses(t *testing.T) { +func TestCancelQueues(t *testing.T) { + // job-states is required by this command, so every case sets it. tests := map[string]struct { - flags []flag - expectError bool + args []string + flags []flag + want []queueCall }{ - "with all-priority-classes flag set": { - flags: []flag{{"all-priority-classes", "true"}, {"job-states", "queued"}}, - expectError: false, + // Omitting priority-classes means all priority classes, which the + // queue API represents as an empty slice. + "without priority-classes": { + args: []string{"queue-a"}, + flags: []flag{{"job-states", "queued"}}, + want: []queueCall{ + { + queue: "queue-a", + priorityClasses: []string{}, + jobStates: []api.JobState{api.JobState_QUEUED}, + pools: []string{}, + }, + }, }, - "without all-priority-classes and without priority-classes": { - flags: []flag{{"job-states", "queued"}}, - expectError: true, + "with a single priority class": { + args: []string{"queue-a"}, + flags: []flag{{"job-states", "queued"}, {"priority-classes", "armada-default"}}, + want: []queueCall{ + { + queue: "queue-a", + priorityClasses: []string{"armada-default"}, + jobStates: []api.JobState{api.JobState_QUEUED}, + pools: []string{}, + }, + }, }, - "without all-priority-classes but with priority-classes": { - flags: []flag{{"priority-classes", "armada-default"}, {"job-states", "queued"}}, - expectError: false, + "with multiple priority classes": { + args: []string{"queue-a"}, + flags: []flag{{"job-states", "queued"}, {"priority-classes", "armada-default,armada-preemptible"}}, + want: []queueCall{ + { + queue: "queue-a", + priorityClasses: []string{"armada-default", "armada-preemptible"}, + jobStates: []api.JobState{api.JobState_QUEUED}, + pools: []string{}, + }, + }, + }, + "with multiple job states and pools": { + args: []string{"queue-a"}, + flags: []flag{{"job-states", "queued,running"}, {"pools", "pool-1"}}, + want: []queueCall{ + { + queue: "queue-a", + priorityClasses: []string{}, + jobStates: []api.JobState{api.JobState_QUEUED, api.JobState_RUNNING}, + pools: []string{"pool-1"}, + }, + }, + }, + "cancels each selected queue": { + args: []string{"queue-a", "queue-b"}, + flags: []flag{{"job-states", "queued"}}, + want: []queueCall{ + { + queue: "queue-a", + priorityClasses: []string{}, + jobStates: []api.JobState{api.JobState_QUEUED}, + pools: []string{}, + }, + { + queue: "queue-b", + priorityClasses: []string{}, + jobStates: []api.JobState{api.JobState_QUEUED}, + pools: []string{}, + }, + }, + }, + // dry-run reports what would happen without calling the API. + "dry-run calls nothing": { + args: []string{"queue-a"}, + flags: []flag{{"job-states", "queued"}, {"dry-run", "true"}}, + want: nil, }, } for name, tc := range tests { t.Run(name, func(t *testing.T) { - cmd := cancelQueueCmd() - cmd.PreRunE = func(cmd *cobra.Command, args []string) error { - if err := cmd.MarkFlagRequired("job-states"); err != nil { - return err - } - all, err := cmd.Flags().GetBool("all-priority-classes") - if err != nil { - return err - } - if !all { - if err := cmd.MarkFlagRequired("priority-classes"); err != nil { - return err - } + a := armadactl.New() + cmd := cancelQueueCmd(a) + + var got []queueCall + withFakeAPIs(t, a, cmd, func() { + a.Params.QueueAPI.GetAll = func() ([]*api.Queue, error) { return testQueues(), nil } + a.Params.QueueAPI.Cancel = func(queue string, priorityClasses []string, jobStates []api.JobState, pools []string) error { + got = append(got, queueCall{queue, priorityClasses, jobStates, pools}) + return nil } - return nil - } - cmd.RunE = func(cmd *cobra.Command, args []string) error { - return nil - } - cmd.SetArgs([]string{"test-queue"}) + }) + + cmd.SetArgs(tc.args) for _, f := range tc.flags { require.NoError(t, cmd.Flags().Set(f.name, f.value)) } - err := cmd.Execute() - if tc.expectError { - require.Error(t, err) - } else { - require.NoError(t, err) - } + + require.NoError(t, cmd.Execute()) + require.Equal(t, tc.want, got) }) } } + +func TestCancelQueuesRequiresQueueSelection(t *testing.T) { + // Guards against accidentally cancelling every queue: selection must be + // narrowed by name or by label. + a := armadactl.New() + cmd := cancelQueueCmd(a) + + called := false + withFakeAPIs(t, a, cmd, func() { + a.Params.QueueAPI.GetAll = func() ([]*api.Queue, error) { return testQueues(), nil } + a.Params.QueueAPI.Cancel = func(queue string, priorityClasses []string, jobStates []api.JobState, pools []string) error { + called = true + return nil + } + }) + + // Must be non-nil: cobra falls back to os.Args[1:] when args are nil. + cmd.SetArgs([]string{}) + require.NoError(t, cmd.Flags().Set("job-states", "queued")) + cmd.SilenceUsage = true + + require.Error(t, cmd.Execute()) + require.False(t, called, "no queue should be cancelled without a selection") +} diff --git a/cmd/armadactl/cmd/preempt.go b/cmd/armadactl/cmd/preempt.go index 8ed06d490cb..83170d8c710 100644 --- a/cmd/armadactl/cmd/preempt.go +++ b/cmd/armadactl/cmd/preempt.go @@ -18,9 +18,9 @@ func preemptCmd() *cobra.Command { } cmd.AddCommand( preemptJobCmd(), - preemptExecutorCmd(), - preemptNodeCmd(), - preemptQueuesCmd(), + preemptExecutorCmd(armadactl.New()), + preemptNodeCmd(armadactl.New()), + preemptQueuesCmd(armadactl.New()), ) return cmd } @@ -46,23 +46,13 @@ func preemptJobCmd() *cobra.Command { return cmd } -func preemptExecutorCmd() *cobra.Command { - a := armadactl.New() +func preemptExecutorCmd(a *armadactl.App) *cobra.Command { cmd := &cobra.Command{ Use: "executor ", Short: "Preempts jobs on executor.", Long: `Preempts jobs on executor with provided executor name, priority classes, queues, and pools.`, Args: cobra.ExactArgs(1), PreRunE: func(cmd *cobra.Command, args []string) error { - all, err := cmd.Flags().GetBool("all-priority-classes") - if err != nil { - return fmt.Errorf("error reading all-priority-classes flag: %s", err) - } - if !all { - if err := cmd.MarkFlagRequired("priority-classes"); err != nil { - return fmt.Errorf("error marking priority-class flag as required: %s", err) - } - } return initParams(cmd, a.Params) }, RunE: func(cmd *cobra.Command, args []string) error { @@ -73,11 +63,6 @@ func preemptExecutorCmd() *cobra.Command { return fmt.Errorf("error reading priority-class selection: %s", err) } - allPriorityClasses, _ := cmd.Flags().GetBool("all-priority-classes") - if allPriorityClasses { - priorityClasses = nil - } - queues, err := cmd.Flags().GetStringSlice("queues") if err != nil { return fmt.Errorf("error reading queue selection: %s", err) @@ -101,13 +86,7 @@ func preemptExecutorCmd() *cobra.Command { "priority-classes", "p", []string{}, - "Preempt jobs on executor matching the specified priority classes. Provided priority classes should be comma separated, as in the following example: armada-default,armada-preemptible.", - ) - cmd.Flags().BoolP( - "all-priority-classes", - "a", - false, - "Preempt jobs on executor for all priority classes.", + "Preempt jobs on executor matching the specified priority classes, comma separated (e.g. armada-default,armada-preemptible). If no priority classes are provided, jobs across all priority classes will be preempted.", ) cmd.Flags().StringSlice( "pools", @@ -117,23 +96,13 @@ func preemptExecutorCmd() *cobra.Command { return cmd } -func preemptNodeCmd() *cobra.Command { - a := armadactl.New() +func preemptNodeCmd(a *armadactl.App) *cobra.Command { cmd := &cobra.Command{ Use: "node ", Short: "Preempts jobs on node for specified executor.", Long: `Preempts jobs on node for specified executor with provided node name, executor name, priority classes and queues.`, Args: cobra.ExactArgs(1), PreRunE: func(cmd *cobra.Command, args []string) error { - all, err := cmd.Flags().GetBool("all-priority-classes") - if err != nil { - return fmt.Errorf("error reading all-priority-classes flag: %s", err) - } - if !all { - if err := cmd.MarkFlagRequired("priority-classes"); err != nil { - return fmt.Errorf("error marking priority-class flag as required: %s", err) - } - } if err := cmd.MarkFlagRequired("executor"); err != nil { return fmt.Errorf("error marking executor flag as required: %s", err) } @@ -147,11 +116,6 @@ func preemptNodeCmd() *cobra.Command { return fmt.Errorf("error reading priority-class selection: %s", err) } - allPriorityClasses, _ := cmd.Flags().GetBool("all-priority-classes") - if allPriorityClasses { - priorityClasses = nil - } - queues, err := cmd.Flags().GetStringSlice("queues") if err != nil { return fmt.Errorf("error reading queue selection: %s", err) @@ -175,13 +139,7 @@ func preemptNodeCmd() *cobra.Command { "priority-classes", "p", []string{}, - "Preempt jobs on node for specified executor matching the specified priority classes. Provided priority classes should be comma separated, as in the following example: armada-default,armada-preemptible.", - ) - cmd.Flags().BoolP( - "all-priority-classes", - "a", - false, - "Preempt jobs on node for specified executor for all priority classes.", + "Preempt jobs on node for specified executor matching the specified priority classes, comma separated (e.g. armada-default,armada-preemptible). If no priority classes are provided, jobs across all priority classes will be preempted.", ) cmd.Flags().StringP( "executor", @@ -192,23 +150,13 @@ func preemptNodeCmd() *cobra.Command { return cmd } -func preemptQueuesCmd() *cobra.Command { - a := armadactl.New() +func preemptQueuesCmd(a *armadactl.App) *cobra.Command { cmd := &cobra.Command{ Use: "queues ...", Short: "Preempts jobs on queues.", Long: `Preempts jobs on selected queues in specified priority classes and pools. Allows selecting of queues by label or name, one of which must be provided. All flags with multiple values must be comma separated.`, Aliases: []string{"queue"}, PreRunE: func(cmd *cobra.Command, args []string) error { - all, err := cmd.Flags().GetBool("all-priority-classes") - if err != nil { - return fmt.Errorf("error reading all-priority-classes flag: %s", err) - } - if !all { - if err := cmd.MarkFlagRequired("priority-classes"); err != nil { - return fmt.Errorf("error marking priority-class flag as required: %s", err) - } - } return initParams(cmd, a.Params) }, RunE: func(cmd *cobra.Command, queues []string) error { @@ -237,11 +185,6 @@ func preemptQueuesCmd() *cobra.Command { return fmt.Errorf("error reading priority-classes flag: %s", err) } - allPriorityClasses, _ := cmd.Flags().GetBool("all-priority-classes") - if allPriorityClasses { - priorityClasses = nil - } - pools, err := cmd.Flags().GetStringSlice("pools") if err != nil { return fmt.Errorf("error reading pools flag: %s", err) @@ -271,13 +214,7 @@ func preemptQueuesCmd() *cobra.Command { "priority-classes", "p", []string{}, - "Jobs matching the provided priority classes will be preempted.", - ) - cmd.Flags().BoolP( - "all-priority-classes", - "a", - false, - "Preempt jobs in all priority classes. Cannot be used with the priority-classes flag.", + "Jobs matching the provided priority classes will be preempted. If no priority classes are provided, jobs across all priority classes will be preempted.", ) cmd.Flags().StringSliceP( "selector", diff --git a/cmd/armadactl/cmd/preempt_test.go b/cmd/armadactl/cmd/preempt_test.go index 9aa2c09e119..19194189d10 100644 --- a/cmd/armadactl/cmd/preempt_test.go +++ b/cmd/armadactl/cmd/preempt_test.go @@ -8,214 +8,291 @@ import ( "github.com/stretchr/testify/require" "github.com/armadaproject/armada/internal/armadactl" + "github.com/armadaproject/armada/pkg/api" ) -func TestPreemptQueue(t *testing.T) { - tests := map[string]struct { - Flags []flag - selectors []string - priorityClasses []string - inverse bool - onlyCordoned bool - dryRun bool - }{ - "default flags": {nil, []string{}, []string{}, false, false, false}, - "valid selectors": {[]flag{{"selectors", "armadaproject.io/priority=high,armadaproject.io/category=critical"}}, []string{"armadaproject.io/priority=high", "armadaproject.io/category=critical"}, []string{}, false, false, false}, - "valid priority-classes 1": {[]flag{{"priority-classes", "armada-default"}}, []string{}, []string{"armada-default"}, false, false, false}, - "valid priority-classes 2": {[]flag{{"priority-classes", "armada-default,armada-preemptible"}}, []string{}, []string{"armada-default", "armada-preemptible"}, false, false, false}, - "valid multiple flags": { - []flag{{"selectors", "armadaproject.io/priority=high,armadaproject.io/category=critical"}, {"priority-classes", "armada-default,armada-preemptible"}}, - []string{"armadaproject.io/priority=high", "armadaproject.io/category=critical"}, - []string{"armada-default", "armada-preemptible"}, - true, true, true, - }, - } - for name, test := range tests { - t.Run(name, func(t *testing.T) { - a := armadactl.New() - cmd := preemptQueuesCmd() - - cmd.PreRunE = func(cmd *cobra.Command, args []string) error { - a.Out = io.Discard - - if len(test.selectors) > 0 { - selectorsFlag, err := cmd.Flags().GetString("selectors") - require.Error(t, err) - require.Equal(t, test.selectors, selectorsFlag) - } - if len(test.priorityClasses) > 0 { - priorityClassesFlag, err := cmd.Flags().GetString("priority-classes") - require.Error(t, err) - require.Equal(t, test.priorityClasses, priorityClassesFlag) - } +// executorCall records the arguments armadactl passes to the executor API. +type executorCall struct { + executor string + queues []string + priorityClasses []string + pools []string +} - inverseValue, err := cmd.Flags().GetBool("inverse") - require.NoError(t, err) - require.Equal(t, test, inverseValue) +// nodeCall records the arguments armadactl passes to the node API. +type nodeCall struct { + node string + executor string + queues []string + priorityClasses []string +} - onlyCordonedValue, err := cmd.Flags().GetBool("only-cordoned") - require.NoError(t, err) - require.Equal(t, test, onlyCordonedValue) +// queueCall records the arguments armadactl passes to the queue API. +type queueCall struct { + queue string + priorityClasses []string + jobStates []api.JobState + pools []string +} - dryRunValue, err := cmd.Flags().GetBool("dry-run") - require.NoError(t, err) - require.Equal(t, test, dryRunValue) +// testQueues is the set of queues the faked GetAll returns. Commands that are +// not narrowed by queue expand to all of them. +func testQueues() []*api.Queue { + return []*api.Queue{{Name: "queue-a"}, {Name: "queue-b"}} +} - return nil - } - }) +// withFakeAPIs runs the command's real PreRunE and then replaces the API +// functions that initParams just installed, so the command executes its real +// RunE all the way to the API boundary without making network calls. +// +// installFakes must overwrite the APIs *after* initParams has run, since +// initParams assigns every function pointer in Params. +func withFakeAPIs(t *testing.T, a *armadactl.App, cmd *cobra.Command, installFakes func()) { + t.Helper() + a.Out = io.Discard + realPreRunE := cmd.PreRunE + require.NotNil(t, realPreRunE, "expected the command to define a PreRunE") + cmd.PreRunE = func(cmd *cobra.Command, args []string) error { + if err := realPreRunE(cmd, args); err != nil { + return err + } + installFakes() + return nil } } -func TestPreemptExecutorAllPriorityClasses(t *testing.T) { +func TestPreemptExecutor(t *testing.T) { tests := map[string]struct { - flags []flag - expectError bool + flags []flag + want executorCall }{ - "with all-priority-classes flag set": { - flags: []flag{{"all-priority-classes", "true"}}, - expectError: false, + // Omitting priority-classes means all priority classes, which the + // executor API represents as an empty slice. An unnarrowed queue + // selection expands to every queue. + "without priority-classes": { + flags: nil, + want: executorCall{ + executor: "test-executor", + queues: []string{"queue-a", "queue-b"}, + priorityClasses: []string{}, + pools: []string{}, + }, + }, + "with a single priority class": { + flags: []flag{{"priority-classes", "armada-default"}}, + want: executorCall{ + executor: "test-executor", + queues: []string{"queue-a", "queue-b"}, + priorityClasses: []string{"armada-default"}, + pools: []string{}, + }, }, - "without all-priority-classes and without priority-classes": { - flags: nil, - expectError: true, + "with multiple priority classes": { + flags: []flag{{"priority-classes", "armada-default,armada-preemptible"}}, + want: executorCall{ + executor: "test-executor", + queues: []string{"queue-a", "queue-b"}, + priorityClasses: []string{"armada-default", "armada-preemptible"}, + pools: []string{}, + }, }, - "without all-priority-classes but with priority-classes": { - flags: []flag{{"priority-classes", "armada-default"}}, - expectError: false, + "with queues and pools": { + flags: []flag{{"queues", "queue-a"}, {"pools", "pool-1,pool-2"}}, + want: executorCall{ + executor: "test-executor", + queues: []string{"queue-a"}, + priorityClasses: []string{}, + pools: []string{"pool-1", "pool-2"}, + }, }, } for name, tc := range tests { t.Run(name, func(t *testing.T) { - cmd := preemptExecutorCmd() - cmd.PreRunE = func(cmd *cobra.Command, args []string) error { - all, err := cmd.Flags().GetBool("all-priority-classes") - if err != nil { - return err - } - if !all { - if err := cmd.MarkFlagRequired("priority-classes"); err != nil { - return err - } + a := armadactl.New() + cmd := preemptExecutorCmd(a) + + var got []executorCall + withFakeAPIs(t, a, cmd, func() { + a.Params.QueueAPI.GetAll = func() ([]*api.Queue, error) { return testQueues(), nil } + a.Params.ExecutorAPI.PreemptOnExecutor = func(executor string, queues, priorityClasses, pools []string) error { + got = append(got, executorCall{executor, queues, priorityClasses, pools}) + return nil } - return nil - } - cmd.RunE = func(cmd *cobra.Command, args []string) error { - return nil - } + }) + cmd.SetArgs([]string{"test-executor"}) for _, f := range tc.flags { require.NoError(t, cmd.Flags().Set(f.name, f.value)) } - err := cmd.Execute() - if tc.expectError { - require.Error(t, err) - } else { - require.NoError(t, err) - } + + require.NoError(t, cmd.Execute()) + require.Equal(t, []executorCall{tc.want}, got) }) } } -func TestPreemptNodeAllPriorityClasses(t *testing.T) { +func TestPreemptNode(t *testing.T) { tests := map[string]struct { - flags []flag - expectError bool + flags []flag + want nodeCall }{ - "with all-priority-classes flag set": { - flags: []flag{{"all-priority-classes", "true"}, {"executor", "test-exec"}}, - expectError: false, + // Omitting priority-classes means all priority classes, which the node + // API represents as an empty slice. + "without priority-classes": { + flags: []flag{{"executor", "test-executor"}}, + want: nodeCall{ + node: "test-node", + executor: "test-executor", + queues: []string{"queue-a", "queue-b"}, + priorityClasses: []string{}, + }, + }, + "with a single priority class": { + flags: []flag{{"executor", "test-executor"}, {"priority-classes", "armada-default"}}, + want: nodeCall{ + node: "test-node", + executor: "test-executor", + queues: []string{"queue-a", "queue-b"}, + priorityClasses: []string{"armada-default"}, + }, }, - "without all-priority-classes and without priority-classes": { - flags: []flag{{"executor", "test-exec"}}, - expectError: true, + "with multiple priority classes": { + flags: []flag{{"executor", "test-executor"}, {"priority-classes", "armada-default,armada-preemptible"}}, + want: nodeCall{ + node: "test-node", + executor: "test-executor", + queues: []string{"queue-a", "queue-b"}, + priorityClasses: []string{"armada-default", "armada-preemptible"}, + }, }, - "without all-priority-classes but with priority-classes": { - flags: []flag{{"priority-classes", "armada-default"}, {"executor", "test-exec"}}, - expectError: false, + "with queues": { + flags: []flag{{"executor", "test-executor"}, {"queues", "queue-a"}}, + want: nodeCall{ + node: "test-node", + executor: "test-executor", + queues: []string{"queue-a"}, + priorityClasses: []string{}, + }, }, } for name, tc := range tests { t.Run(name, func(t *testing.T) { - cmd := preemptNodeCmd() - cmd.PreRunE = func(cmd *cobra.Command, args []string) error { - all, err := cmd.Flags().GetBool("all-priority-classes") - if err != nil { - return err - } - if !all { - if err := cmd.MarkFlagRequired("priority-classes"); err != nil { - return err - } - } - if err := cmd.MarkFlagRequired("executor"); err != nil { - return err + a := armadactl.New() + cmd := preemptNodeCmd(a) + + var got []nodeCall + withFakeAPIs(t, a, cmd, func() { + a.Params.QueueAPI.GetAll = func() ([]*api.Queue, error) { return testQueues(), nil } + a.Params.NodeAPI.PreemptOnNode = func(node, executor string, queues, priorityClasses []string) error { + got = append(got, nodeCall{node, executor, queues, priorityClasses}) + return nil } - return nil - } - cmd.RunE = func(cmd *cobra.Command, args []string) error { - return nil - } + }) + cmd.SetArgs([]string{"test-node"}) for _, f := range tc.flags { require.NoError(t, cmd.Flags().Set(f.name, f.value)) } - err := cmd.Execute() - if tc.expectError { - require.Error(t, err) - } else { - require.NoError(t, err) - } + + require.NoError(t, cmd.Execute()) + require.Equal(t, []nodeCall{tc.want}, got) }) } } -func TestPreemptQueuesAllPriorityClasses(t *testing.T) { +func TestPreemptQueues(t *testing.T) { tests := map[string]struct { - flags []flag - expectError bool + args []string + flags []flag + want []queueCall }{ - "with all-priority-classes flag set": { - flags: []flag{{"all-priority-classes", "true"}}, - expectError: false, + // Omitting priority-classes means all priority classes, which the + // queue API represents as an empty slice. + "without priority-classes": { + args: []string{"queue-a"}, + flags: nil, + want: []queueCall{ + {queue: "queue-a", priorityClasses: []string{}, pools: []string{}}, + }, + }, + "with a single priority class": { + args: []string{"queue-a"}, + flags: []flag{{"priority-classes", "armada-default"}}, + want: []queueCall{ + {queue: "queue-a", priorityClasses: []string{"armada-default"}, pools: []string{}}, + }, }, - "without all-priority-classes and without priority-classes": { - flags: nil, - expectError: true, + "with multiple priority classes": { + args: []string{"queue-a"}, + flags: []flag{{"priority-classes", "armada-default,armada-preemptible"}}, + want: []queueCall{ + {queue: "queue-a", priorityClasses: []string{"armada-default", "armada-preemptible"}, pools: []string{}}, + }, }, - "without all-priority-classes but with priority-classes": { - flags: []flag{{"priority-classes", "armada-default"}}, - expectError: false, + "with pools": { + args: []string{"queue-a"}, + flags: []flag{{"pools", "pool-1"}}, + want: []queueCall{ + {queue: "queue-a", priorityClasses: []string{}, pools: []string{"pool-1"}}, + }, + }, + "preempts each selected queue": { + args: []string{"queue-a", "queue-b"}, + flags: nil, + want: []queueCall{ + {queue: "queue-a", priorityClasses: []string{}, pools: []string{}}, + {queue: "queue-b", priorityClasses: []string{}, pools: []string{}}, + }, + }, + // dry-run reports what would happen without calling the API. + "dry-run calls nothing": { + args: []string{"queue-a"}, + flags: []flag{{"dry-run", "true"}}, + want: nil, }, } for name, tc := range tests { t.Run(name, func(t *testing.T) { - cmd := preemptQueuesCmd() - cmd.PreRunE = func(cmd *cobra.Command, args []string) error { - all, err := cmd.Flags().GetBool("all-priority-classes") - if err != nil { - return err - } - if !all { - if err := cmd.MarkFlagRequired("priority-classes"); err != nil { - return err - } + a := armadactl.New() + cmd := preemptQueuesCmd(a) + + var got []queueCall + withFakeAPIs(t, a, cmd, func() { + a.Params.QueueAPI.GetAll = func() ([]*api.Queue, error) { return testQueues(), nil } + a.Params.QueueAPI.Preempt = func(queue string, priorityClasses, pools []string) error { + got = append(got, queueCall{queue: queue, priorityClasses: priorityClasses, pools: pools}) + return nil } - return nil - } - cmd.RunE = func(cmd *cobra.Command, args []string) error { - return nil - } - cmd.SetArgs([]string{"test-queue"}) + }) + + cmd.SetArgs(tc.args) for _, f := range tc.flags { require.NoError(t, cmd.Flags().Set(f.name, f.value)) } - err := cmd.Execute() - if tc.expectError { - require.Error(t, err) - } else { - require.NoError(t, err) - } + + require.NoError(t, cmd.Execute()) + require.Equal(t, tc.want, got) }) } } + +func TestPreemptQueuesRequiresQueueSelection(t *testing.T) { + a := armadactl.New() + cmd := preemptQueuesCmd(a) + + called := false + withFakeAPIs(t, a, cmd, func() { + a.Params.QueueAPI.GetAll = func() ([]*api.Queue, error) { return testQueues(), nil } + a.Params.QueueAPI.Preempt = func(queue string, priorityClasses, pools []string) error { + called = true + return nil + } + }) + + // Must be non-nil: cobra falls back to os.Args[1:] when args are nil. + cmd.SetArgs([]string{}) + cmd.SilenceUsage = true + + require.Error(t, cmd.Execute()) + require.False(t, called, "no queue should be preempted without a selection") +}