Scheduler: tidy queue_scheduler_test (#4117)

Signed-off-by: Robert Smith <[email protected]>
armadaproject · Jan 3, 2025 · 6899a5f · 6899a5f
1 parent 63b9202
commit 6899a5f
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 65 deletions.
diff --git a/internal/scheduler/internaltypes/resource_list_map_util.go b/internal/scheduler/internaltypes/resource_list_map_util.go
@@ -2,8 +2,6 @@ package internaltypes
 
 import (
 	"strings"
-
-	"github.com/armadaproject/armada/internal/scheduler/schedulerobjects"
 )
 
 func RlMapToString(m map[string]ResourceList) string {
@@ -40,14 +38,6 @@ func RlMapHasNegativeValues(m map[string]ResourceList) bool {
 	return false
 }
 
-func RlMapFromJobSchedulerObjects(m schedulerobjects.QuantityByTAndResourceType[string], rlFactory *ResourceListFactory) map[string]ResourceList {
-	result := map[string]ResourceList{}
-	for k, v := range m {
-		result[k] = rlFactory.FromJobResourceListIgnoreUnknown(v.Resources)
-	}
-	return result
-}
-
 func RlMapRemoveZeros(m map[string]ResourceList) map[string]ResourceList {
 	result := map[string]ResourceList{}
 	for k, v := range m {

diff --git a/internal/scheduler/internaltypes/resource_list_map_util_test.go b/internal/scheduler/internaltypes/resource_list_map_util_test.go
@@ -4,9 +4,6 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"k8s.io/apimachinery/pkg/api/resource"
-
-	"github.com/armadaproject/armada/internal/scheduler/schedulerobjects"
 )
 
 func TestRlMapSumValues(t *testing.T) {
@@ -36,35 +33,6 @@ func TestRlMapHasNegativeValues(t *testing.T) {
 	assert.False(t, RlMapHasNegativeValues(testMapEmpty(factory)))
 }
 
-func TestRlMapFromJobSchedulerObjects(t *testing.T) {
-	factory := testFactory()
-
-	input := make(schedulerobjects.QuantityByTAndResourceType[string])
-	input.AddResourceList("priorityClass1",
-		schedulerobjects.ResourceList{
-			Resources: map[string]resource.Quantity{
-				"cpu":    resource.MustParse("1"),
-				"memory": resource.MustParse("1Ki"),
-			},
-		},
-	)
-	input.AddResourceList("priorityClass2",
-		schedulerobjects.ResourceList{
-			Resources: map[string]resource.Quantity{
-				"cpu":    resource.MustParse("2"),
-				"memory": resource.MustParse("2Ki"),
-			},
-		},
-	)
-
-	expected := map[string]ResourceList{
-		"priorityClass1": testResourceList(factory, "1", "1Ki"),
-		"priorityClass2": testResourceList(factory, "2", "2Ki"),
-	}
-
-	assert.Equal(t, expected, RlMapFromJobSchedulerObjects(input, factory))
-}
-
 func TestRlMapRemoveZeros(t *testing.T) {
 	factory := testFactory()
 

diff --git a/internal/scheduler/scheduling/queue_scheduler_test.go b/internal/scheduler/scheduling/queue_scheduler_test.go
@@ -11,7 +11,6 @@ import (
 	"golang.org/x/exp/slices"
 	"golang.org/x/time/rate"
 	v1 "k8s.io/api/core/v1"
-	"k8s.io/apimachinery/pkg/api/resource"
 
 	"github.com/armadaproject/armada/internal/common/armadacontext"
 	armadaslices "github.com/armadaproject/armada/internal/common/slices"
@@ -20,7 +19,6 @@ import (
 	"github.com/armadaproject/armada/internal/scheduler/internaltypes"
 	"github.com/armadaproject/armada/internal/scheduler/jobdb"
 	"github.com/armadaproject/armada/internal/scheduler/nodedb"
-	"github.com/armadaproject/armada/internal/scheduler/schedulerobjects"
 	schedulerconstraints "github.com/armadaproject/armada/internal/scheduler/scheduling/constraints"
 	"github.com/armadaproject/armada/internal/scheduler/scheduling/context"
 	"github.com/armadaproject/armada/internal/scheduler/scheduling/fairness"
@@ -32,13 +30,10 @@ import (
 func TestQueueScheduler(t *testing.T) {
 	tests := map[string]struct {
 		SchedulingConfig configuration.SchedulingConfig
-		// Total resources across all clusters.
-		// Set to the total resources across all nodes if not provided.
-		TotalResources schedulerobjects.ResourceList
 		// Queues
 		Queues []*api.Queue
 		// Initial resource usage for all queues.
-		InitialAllocatedByQueueAndPriorityClass map[string]schedulerobjects.QuantityByTAndResourceType[string]
+		InitialAllocatedByQueueAndPriorityClass map[string]map[string]internaltypes.ResourceList
 		// Nodes to be considered by the scheduler.
 		Nodes []*internaltypes.Node
 		// Jobs to try scheduling.
@@ -263,13 +258,9 @@ func TestQueueScheduler(t *testing.T) {
 				testfixtures.N1Cpu4GiJobs("B", testfixtures.PriorityClass0, 32),
 			),
 			Queues: []*api.Queue{{Name: "A", PriorityFactor: 1.0}, {Name: "B", PriorityFactor: 1.0}},
-			InitialAllocatedByQueueAndPriorityClass: map[string]schedulerobjects.QuantityByTAndResourceType[string]{
+			InitialAllocatedByQueueAndPriorityClass: map[string]map[string]internaltypes.ResourceList{
 				"A": {
-					testfixtures.PriorityClass0: schedulerobjects.ResourceList{
-						Resources: map[string]resource.Quantity{
-							"cpu": resource.MustParse("100"),
-						},
-					},
+					testfixtures.PriorityClass0: testfixtures.Cpu("100"),
 				},
 			},
 			ExpectedScheduledIndices: testfixtures.IntRange(32, 63),
@@ -476,12 +467,8 @@ func TestQueueScheduler(t *testing.T) {
 				require.NoError(t, err)
 			}
 			txn.Commit()
-			if tc.TotalResources.Resources == nil {
-				// Default to NodeDb total.
-				tc.TotalResources = schedulerobjects.ResourceList{
-					Resources: nodeDb.TotalKubernetesResources().ToMap(),
-				}
-			}
+
+			totalResources := nodeDb.TotalKubernetesResources()
 
 			queueNameToQueue := map[string]*api.Queue{}
 			for _, q := range tc.Queues {
@@ -500,7 +487,6 @@ func TestQueueScheduler(t *testing.T) {
 				context.JobSchedulingContextsFromJobs(tc.Jobs),
 			)
 
-			totalResources := testfixtures.TestResourceListFactory.FromJobResourceListIgnoreUnknown(tc.TotalResources.Resources)
 			fairnessCostProvider, err := fairness.NewDominantResourceFairness(
 				totalResources,
 				tc.SchedulingConfig,
@@ -519,10 +505,7 @@ func TestQueueScheduler(t *testing.T) {
 				weight := 1.0 / float64(q.PriorityFactor)
 				err := sctx.AddQueueSchedulingContext(
 					q.Name, weight,
-					internaltypes.RlMapFromJobSchedulerObjects(
-						tc.InitialAllocatedByQueueAndPriorityClass[q.Name],
-						testfixtures.TestResourceListFactory,
-					),
+					tc.InitialAllocatedByQueueAndPriorityClass[q.Name],
 					internaltypes.ResourceList{},
 					internaltypes.ResourceList{},
 					rate.NewLimiter(