Scheduler: refactor: move scheduling tests to internaltypes (#4113)

Signed-off-by: Robert Smith <robertdavidsmith@yahoo.com>
armadaproject · Jan 3, 2025 · e1e91a9 · e1e91a9
1 parent df5401c
commit e1e91a9
Show file tree

Hide file tree

Showing 11 changed files with 388 additions and 587 deletions.
diff --git a/internal/scheduler/internaltypes/resource_list_map_util.go b/internal/scheduler/internaltypes/resource_list_map_util.go
@@ -63,6 +63,7 @@ func NewAllocatableByPriorityAndResourceType(priorities []int32, rl ResourceList
 	for _, priority := range priorities {
 		result[priority] = rl
 	}
+	result[EvictedPriority] = rl
 	return result
 }
 

diff --git a/internal/scheduler/internaltypes/resource_list_map_util_test.go b/internal/scheduler/internaltypes/resource_list_map_util_test.go
@@ -83,9 +83,10 @@ func TestNewAllocatableByPriorityAndResourceType(t *testing.T) {
 	rl := testResourceList(factory, "2", "2Ki")
 
 	result := NewAllocatableByPriorityAndResourceType([]int32{1, 2}, rl)
-	assert.Equal(t, 2, len(result))
+	assert.Equal(t, 3, len(result))
 	assert.Equal(t, int64(2000), result[1].GetByNameZeroIfMissing("cpu"))
 	assert.Equal(t, int64(2000), result[2].GetByNameZeroIfMissing("cpu"))
+	assert.Equal(t, int64(2000), result[EvictedPriority].GetByNameZeroIfMissing("cpu"))
 }
 
 func TestMarkAllocated(t *testing.T) {

diff --git a/internal/scheduler/nodedb/nodedb_test.go b/internal/scheduler/nodedb/nodedb_test.go
@@ -833,7 +833,7 @@ func BenchmarkScheduleMany100CpuNodes1CpuUnused(b *testing.B) {
 		b,
 		testfixtures.ItWithUsedResourcesNodes(
 			0,
-			cpu("31"),
+			testfixtures.Cpu("31"),
 			testfixtures.ItN32CpuNodes(100, testfixtures.TestPriorities),
 		),
 		testfixtures.N1Cpu4GiJobs("A", testfixtures.PriorityClass0, 100),
@@ -845,7 +845,7 @@ func BenchmarkScheduleMany1000CpuNodes1CpuUnused(b *testing.B) {
 		b,
 		testfixtures.ItWithUsedResourcesNodes(
 			0,
-			cpu("31"),
+			testfixtures.Cpu("31"),
 			testfixtures.ItN32CpuNodes(1000, testfixtures.TestPriorities),
 		),
 		testfixtures.N1Cpu4GiJobs("A", testfixtures.PriorityClass0, 1000),
@@ -857,7 +857,7 @@ func BenchmarkScheduleMany10000CpuNodes1CpuUnused(b *testing.B) {
 		b,
 		testfixtures.ItWithUsedResourcesNodes(
 			0,
-			cpu("31"),
+			testfixtures.Cpu("31"),
 			testfixtures.ItN32CpuNodes(10000, testfixtures.TestPriorities),
 		),
 		testfixtures.N1Cpu4GiJobs("A", testfixtures.PriorityClass0, 10000),