diff --git a/internal/scheduler/scheduling/preemption_description_test.go b/internal/scheduler/scheduling/preemption_description_test.go index 92f8dd60692..0d70f3dad0a 100644 --- a/internal/scheduler/scheduling/preemption_description_test.go +++ b/internal/scheduler/scheduling/preemption_description_test.go @@ -14,137 +14,98 @@ import ( "github.com/armadaproject/armada/internal/server/configuration" ) -func TestPopulatePreemptionDescriptions_UnknownCause(t *testing.T) { - scheduledJobContexts := []*context.JobSchedulingContext{} - expectedScheduleJobContexts := []*context.JobSchedulingContext{} - - preemptedJobContexts := []*context.JobSchedulingContext{ - { - JobId: "job-1", - AssignedNodeId: "node-1", - Job: makeJob(t, "job-1", false), - }, - } - expectedPreemptedJobContexts := []*context.JobSchedulingContext{ - { - JobId: "job-1", - AssignedNodeId: "node-1", - Job: makeJob(t, "job-1", false), - PreemptionDescription: unknownPreemptionCause, - }, - } - - PopulatePreemptionDescriptions(preemptedJobContexts, scheduledJobContexts) - assert.Equal(t, expectedScheduleJobContexts, scheduledJobContexts) - assert.Equal(t, expectedPreemptedJobContexts, preemptedJobContexts) -} - -func TestPopulatePreemptionDescriptions_UnknownGangCause(t *testing.T) { - scheduledJobContexts := []*context.JobSchedulingContext{} - expectedScheduleJobContexts := []*context.JobSchedulingContext{} - - preemptedJobContexts := []*context.JobSchedulingContext{ - { - JobId: "job-1", - AssignedNodeId: "node-1", - Job: makeJob(t, "job-1", true), - }, - } - expectedPreemptedJobContexts := []*context.JobSchedulingContext{ - { - JobId: "job-1", - AssignedNodeId: "node-1", - Job: makeJob(t, "job-1", true), - PreemptionDescription: unknownGangPreemptionCause, - }, - } - - PopulatePreemptionDescriptions(preemptedJobContexts, scheduledJobContexts) - assert.Equal(t, expectedScheduleJobContexts, scheduledJobContexts) - assert.Equal(t, expectedPreemptedJobContexts, preemptedJobContexts) -} - -func TestPopulatePreemptionDescriptions_PreemptedByUrgencyBasedPreemption_Single(t *testing.T) { +func TestPopulatePreemptionDescriptions(t *testing.T) { scheduledJobContexts := []*context.JobSchedulingContext{ makeJobSchedulingContext("job-2", "node-1", context.ScheduledWithUrgencyBasedPreemption), + makeJobSchedulingContext("job-3", "node-2", context.ScheduledWithUrgencyBasedPreemption), + makeJobSchedulingContext("job-4", "node-2", context.ScheduledWithUrgencyBasedPreemption), + makeJobSchedulingContext("job-5", "node-1", context.ScheduledWithFairSharePreemption), + makeJobSchedulingContext("job-5", "node-2", context.ScheduledWithFairSharePreemption), + makeJobSchedulingContext("job-6", "node-3", context.ScheduledWithFairSharePreemption), } expectedScheduleJobContexts := []*context.JobSchedulingContext{ makeJobSchedulingContext("job-2", "node-1", context.ScheduledWithUrgencyBasedPreemption), - } - - preemptedJobContexts := []*context.JobSchedulingContext{ - { - JobId: "job-1", - AssignedNodeId: "node-1", - }, - } - - expectedPreemptedJobContexts := []*context.JobSchedulingContext{ - { - JobId: "job-1", - AssignedNodeId: "node-1", - PreemptionDescription: fmt.Sprintf(urgencyPreemptionTemplate, "job-2"), + makeJobSchedulingContext("job-3", "node-2", context.ScheduledWithUrgencyBasedPreemption), + makeJobSchedulingContext("job-4", "node-2", context.ScheduledWithUrgencyBasedPreemption), + makeJobSchedulingContext("job-5", "node-1", context.ScheduledWithFairSharePreemption), + makeJobSchedulingContext("job-5", "node-2", context.ScheduledWithFairSharePreemption), + makeJobSchedulingContext("job-6", "node-3", context.ScheduledWithFairSharePreemption), + } + + tests := map[string]struct { + preemptedJobContext *context.JobSchedulingContext + expectedPreemptedJobContext *context.JobSchedulingContext + }{ + "unknown cause - basic job": { + preemptedJobContext: &context.JobSchedulingContext{ + JobId: "job-1", + AssignedNodeId: "node-3", + Job: makeJob(t, "job-1", false), + }, + expectedPreemptedJobContext: &context.JobSchedulingContext{ + JobId: "job-1", + AssignedNodeId: "node-3", + Job: makeJob(t, "job-1", false), + PreemptionDescription: unknownPreemptionCause, + }, }, - } - - PopulatePreemptionDescriptions(preemptedJobContexts, scheduledJobContexts) - assert.Equal(t, expectedScheduleJobContexts, scheduledJobContexts) - assert.Equal(t, expectedPreemptedJobContexts, preemptedJobContexts) -} - -func TestPopulatePreemptionDescriptions_PreemptedByUrgencyBasedPreemption_Multiple(t *testing.T) { - scheduledJobContexts := []*context.JobSchedulingContext{ - makeJobSchedulingContext("job-2", "node-1", context.ScheduledWithUrgencyBasedPreemption), - makeJobSchedulingContext("job-3", "node-1", context.ScheduledWithUrgencyBasedPreemption), - } - expectedScheduleJobContexts := []*context.JobSchedulingContext{ - makeJobSchedulingContext("job-2", "node-1", context.ScheduledWithUrgencyBasedPreemption), - makeJobSchedulingContext("job-3", "node-1", context.ScheduledWithUrgencyBasedPreemption), - } - - preemptedJobContexts := []*context.JobSchedulingContext{ - { - JobId: "job-1", - AssignedNodeId: "node-1", + "unknown cause - gang job": { + preemptedJobContext: &context.JobSchedulingContext{ + JobId: "job-1", + AssignedNodeId: "node-3", + Job: makeJob(t, "job-1", true), + }, + expectedPreemptedJobContext: &context.JobSchedulingContext{ + JobId: "job-1", + AssignedNodeId: "node-3", + Job: makeJob(t, "job-1", true), + PreemptionDescription: unknownGangPreemptionCause, + }, }, - } - - expectedPreemptedJobContexts := []*context.JobSchedulingContext{ - { - JobId: "job-1", - AssignedNodeId: "node-1", - PreemptionDescription: fmt.Sprintf(urgencyPreemptionMultiJobTemplate, "job-2,job-3"), + "urgency preemption - single preempting job": { + preemptedJobContext: &context.JobSchedulingContext{ + JobId: "job-1", + AssignedNodeId: "node-1", + }, + expectedPreemptedJobContext: &context.JobSchedulingContext{ + JobId: "job-1", + AssignedNodeId: "node-1", + PreemptionDescription: fmt.Sprintf(urgencyPreemptionTemplate, "job-2"), + }, }, - } - - PopulatePreemptionDescriptions(preemptedJobContexts, scheduledJobContexts) - assert.Equal(t, expectedScheduleJobContexts, scheduledJobContexts) - assert.Equal(t, expectedPreemptedJobContexts, preemptedJobContexts) -} - -func TestPopulatePreemptionDescriptions_PreemptedByFairsharePreemption(t *testing.T) { - scheduledJobContexts := []*context.JobSchedulingContext{} - expectedScheduleJobContexts := []*context.JobSchedulingContext{} - - preemptedJobContexts := []*context.JobSchedulingContext{ - { - JobId: "job-1", - AssignedNodeId: "node-1", - PreemptingJobId: "job-2", + "urgency preemption - multiple preempting jobs": { + preemptedJobContext: &context.JobSchedulingContext{ + JobId: "job-1", + AssignedNodeId: "node-2", + }, + expectedPreemptedJobContext: &context.JobSchedulingContext{ + JobId: "job-1", + AssignedNodeId: "node-2", + PreemptionDescription: fmt.Sprintf(urgencyPreemptionMultiJobTemplate, "job-3,job-4"), + }, }, - } - expectedPreemptedJobContexts := []*context.JobSchedulingContext{ - { - JobId: "job-1", - AssignedNodeId: "node-1", - PreemptingJobId: "job-2", - PreemptionDescription: fmt.Sprintf(fairSharePreemptionTemplate, "job-2"), + "fairshare": { + preemptedJobContext: &context.JobSchedulingContext{ + JobId: "job-1", + AssignedNodeId: "node-4", + PreemptingJobId: "job-7", + }, + expectedPreemptedJobContext: &context.JobSchedulingContext{ + JobId: "job-1", + AssignedNodeId: "node-4", + PreemptingJobId: "job-7", + PreemptionDescription: fmt.Sprintf(fairSharePreemptionTemplate, "job-7"), + }, }, } - PopulatePreemptionDescriptions(preemptedJobContexts, scheduledJobContexts) - assert.Equal(t, expectedScheduleJobContexts, scheduledJobContexts) - assert.Equal(t, expectedPreemptedJobContexts, preemptedJobContexts) + for name, tc := range tests { + t.Run(name, func(t *testing.T) { + PopulatePreemptionDescriptions([]*context.JobSchedulingContext{tc.preemptedJobContext}, scheduledJobContexts) + assert.Equal(t, expectedScheduleJobContexts, scheduledJobContexts) + assert.Equal(t, []*context.JobSchedulingContext{tc.expectedPreemptedJobContext}, []*context.JobSchedulingContext{tc.preemptedJobContext}) + }) + } } func makeJobSchedulingContext(jobId string, nodeId string, schedulingMethod context.SchedulingType) *context.JobSchedulingContext {