Improve tests

Signed-off-by: JamesMurkin <jamesmurkin@hotmail.com>
armadaproject · Oct 24, 2024 · aeefea4 · aeefea4
1 parent 6d9445b
commit aeefea4
Showing 1 changed file with 80 additions and 119 deletions.
diff --git 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 {