Fix the tests

armadaproject · Jun 28, 2023 · c4c8c25 · c4c8c25
1 parent c8f42b1
commit c4c8c25
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 3 deletions.
diff --git a/internal/scheduler/nodedb/nodedb.go b/internal/scheduler/nodedb/nodedb.go
@@ -562,6 +562,7 @@ func EvictJobsFromNode(priorityClasses map[string]configuration.PriorityClass, j
 		if !jobFilter(job) {
 			continue
 		}
+		evicted = append(evicted, job)
 		if err := evictJobFromNodeInPlace(priorityClasses, job, node); err != nil {
 			return nil, nil, err
 		}

diff --git a/internal/scheduler/nodedb/nodedb_test.go b/internal/scheduler/nodedb/nodedb_test.go
@@ -116,6 +116,7 @@ func TestSelectNodeForPod_NodeIdLabel_Failure(t *testing.T) {
 }
 
 func TestNodeBindingEvictionUnbinding(t *testing.T) {
+	jobFilter := func(job interfaces.LegacySchedulerJob) bool { return true }
 	node := testfixtures.Test8GpuNode(append(testfixtures.TestPriorities, evictedPriority))
 	job := testfixtures.Test1GpuJob("A", testfixtures.PriorityClass0)
 	request := schedulerobjects.ResourceListFromV1ResourceList(job.GetResourceRequirements().Requests)
@@ -130,16 +131,18 @@ func TestNodeBindingEvictionUnbinding(t *testing.T) {
 	unboundMultipleNode, err := UnbindJobsFromNode(testfixtures.TestPriorityClasses, []interfaces.LegacySchedulerJob{job}, boundNode)
 	require.NoError(t, err)
 
-	evictedNode, err := EvictJobFromNode(testfixtures.TestPriorityClasses, job, boundNode)
+	evictedJobs, evictedNode, err := EvictJobsFromNode(testfixtures.TestPriorityClasses, jobFilter, []interfaces.LegacySchedulerJob{job}, boundNode)
 	require.NoError(t, err)
+	// TODO(zuqq): Add a separate test for this.
+	assert.Equal(t, []interfaces.LegacySchedulerJob{job}, evictedJobs)
 
 	evictedUnboundNode, err := UnbindJobFromNode(testfixtures.TestPriorityClasses, job, evictedNode)
 	require.NoError(t, err)
 
 	evictedBoundNode, err := BindJobToNode(testfixtures.TestPriorityClasses, job, evictedNode)
 	require.NoError(t, err)
 
-	_, err = EvictJobFromNode(testfixtures.TestPriorityClasses, job, node)
+	_, _, err = EvictJobsFromNode(testfixtures.TestPriorityClasses, jobFilter, []interfaces.LegacySchedulerJob{job}, node)
 	require.Error(t, err)
 
 	_, err = UnbindJobFromNode(testfixtures.TestPriorityClasses, job, node)
@@ -148,7 +151,7 @@ func TestNodeBindingEvictionUnbinding(t *testing.T) {
 	_, err = BindJobToNode(testfixtures.TestPriorityClasses, job, boundNode)
 	require.Error(t, err)
 
-	_, err = EvictJobFromNode(testfixtures.TestPriorityClasses, job, evictedNode)
+	_, _, err = EvictJobsFromNode(testfixtures.TestPriorityClasses, jobFilter, []interfaces.LegacySchedulerJob{job}, evictedNode)
 	require.Error(t, err)
 
 	assertNodeAccountingEqual(t, node, unboundNode)
@@ -255,6 +258,26 @@ func assertNodeAccountingEqual(t *testing.T, node1, node2 *schedulerobjects.Node
 	return rv
 }
 
+func TestEviction(t *testing.T) {
+	jobFilter := func(job interfaces.LegacySchedulerJob) bool {
+		priorityClassName := job.GetPriorityClassName()
+		priorityClass := testfixtures.TestPriorityClasses[priorityClassName]
+		return priorityClass.Preemptible
+	}
+	node := testfixtures.Test32CpuNode(testfixtures.TestPriorities)
+	job0 := testfixtures.Test1CpuJob("queue-alice", testfixtures.PriorityClass0)
+	job3 := testfixtures.Test1CpuJob("queue-alice", testfixtures.PriorityClass3)
+	var err error
+	node, err = BindJobToNode(testfixtures.TestPriorityClasses, job0, node)
+	require.NoError(t, err)
+	node, err = BindJobToNode(testfixtures.TestPriorityClasses, job3, node)
+	require.NoError(t, err)
+
+	evictedJobs, node, err := EvictJobsFromNode(testfixtures.TestPriorityClasses, jobFilter, []interfaces.LegacySchedulerJob{job0, job3}, node)
+	require.NoError(t, err)
+	assert.Equal(t, []interfaces.LegacySchedulerJob{job0}, evictedJobs)
+}
+
 func TestScheduleIndividually(t *testing.T) {
 	tests := map[string]struct {
 		Nodes         []*schedulerobjects.Node