Scheduler: move more scheduling test code to internaltypes

Signed-off-by: Robert Smith <[email protected]>
armadaproject · Jan 3, 2025 · 83ed5bf · 83ed5bf
1 parent e1e91a9
commit 83ed5bf
Show file tree

Hide file tree

Showing 8 changed files with 166 additions and 252 deletions.
diff --git a/internal/scheduler/internaltypes/testfixtures/testfixtures.go b/internal/scheduler/internaltypes/testfixtures/testfixtures.go
diff --git a/internal/scheduler/metrics_test.go b/internal/scheduler/metrics_test.go
@@ -515,7 +515,13 @@ func createExecutor(clusterName string, nodes ...*schedulerobjects.Node) *schedu
 }
 
 func createNode(nodeType string) *schedulerobjects.Node {
-	node := testfixtures.Test32CpuNode([]int32{})
+	node := testfixtures.TestSchedulerObjectsNode(
+		[]int32{},
+		map[string]resource.Quantity{
+			"cpu":    resource.MustParse("32"),
+			"memory": resource.MustParse("256Gi"),
+		},
+	)
 	node.ReportingNodeType = nodeType
 	node.StateByJobRunId = map[string]schedulerobjects.JobRunState{}
 	node.ResourceUsageByQueueAndPool = []*schedulerobjects.PoolQueueResource{}

diff --git a/internal/scheduler/nodedb/nodedb_test.go b/internal/scheduler/nodedb/nodedb_test.go
@@ -14,7 +14,6 @@ import (
 	"github.com/armadaproject/armada/internal/common/util"
 	schedulerconfig "github.com/armadaproject/armada/internal/scheduler/configuration"
 	"github.com/armadaproject/armada/internal/scheduler/internaltypes"
-	ittestfixtures "github.com/armadaproject/armada/internal/scheduler/internaltypes/testfixtures"
 	"github.com/armadaproject/armada/internal/scheduler/jobdb"
 	"github.com/armadaproject/armada/internal/scheduler/schedulerobjects"
 	"github.com/armadaproject/armada/internal/scheduler/scheduling/context"
@@ -74,7 +73,7 @@ func TestSelectNodeForPod_NodeIdLabel_Success(t *testing.T) {
 	jctxs := context.JobSchedulingContextsFromJobs(jobs)
 	for _, jctx := range jctxs {
 		txn := db.Txn(false)
-		jctx.SetAssignedNode(ittestfixtures.TestSimpleNode(nodeId))
+		jctx.SetAssignedNode(testfixtures.TestSimpleNode(nodeId))
 		node, err := db.SelectNodeForJobWithTxn(txn, jctx)
 		txn.Abort()
 		require.NoError(t, err)
@@ -99,7 +98,7 @@ func TestSelectNodeForPod_NodeIdLabel_Failure(t *testing.T) {
 	jctxs := context.JobSchedulingContextsFromJobs(jobs)
 	for _, jctx := range jctxs {
 		txn := db.Txn(false)
-		jctx.SetAssignedNode(ittestfixtures.TestSimpleNode("non-existent node"))
+		jctx.SetAssignedNode(testfixtures.TestSimpleNode("non-existent node"))
 		node, err := db.SelectNodeForJobWithTxn(txn, jctx)
 		txn.Abort()
 		if !assert.NoError(t, err) {

diff --git a/internal/scheduler/scheduling/context/job_test.go b/internal/scheduler/scheduling/context/job_test.go
@@ -6,7 +6,6 @@ import (
 	"github.com/stretchr/testify/assert"
 
 	"github.com/armadaproject/armada/internal/scheduler/configuration"
-	ittestfixtures "github.com/armadaproject/armada/internal/scheduler/internaltypes/testfixtures"
 	"github.com/armadaproject/armada/internal/scheduler/testfixtures"
 )
 
@@ -23,7 +22,7 @@ func TestJobSchedulingContext_SetAssignedNode(t *testing.T) {
 	assert.Empty(t, jctx.GetAssignedNodeId())
 	assert.Empty(t, jctx.AdditionalNodeSelectors)
 
-	n := ittestfixtures.TestSimpleNode("node1")
+	n := testfixtures.TestSimpleNode("node1")
 	jctx.SetAssignedNode(n)
 	assert.Equal(t, n, jctx.GetAssignedNode())
 	assert.Equal(t, "node1", jctx.GetAssignedNodeId())

diff --git a/internal/scheduler/scheduling/preemption_description_test.go b/internal/scheduler/scheduling/preemption_description_test.go
@@ -7,7 +7,6 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
-	ittestfixtures "github.com/armadaproject/armada/internal/scheduler/internaltypes/testfixtures"
 	"github.com/armadaproject/armada/internal/scheduler/jobdb"
 	"github.com/armadaproject/armada/internal/scheduler/schedulerobjects"
 	"github.com/armadaproject/armada/internal/scheduler/scheduling/context"
@@ -40,60 +39,60 @@ func TestPopulatePreemptionDescriptions(t *testing.T) {
 		"unknown cause - basic job": {
 			preemptedJobContext: &context.JobSchedulingContext{
 				JobId:        "job-1",
-				AssignedNode: ittestfixtures.TestSimpleNode("node-3"),
+				AssignedNode: testfixtures.TestSimpleNode("node-3"),
 				Job:          makeJob(t, "job-1", false),
 			},
 			expectedPreemptedJobContext: &context.JobSchedulingContext{
 				JobId:                 "job-1",
-				AssignedNode:          ittestfixtures.TestSimpleNode("node-3"),
+				AssignedNode:          testfixtures.TestSimpleNode("node-3"),
 				Job:                   makeJob(t, "job-1", false),
-				PreemptionDescription: fmt.Sprintf(unknownPreemptionCause, ittestfixtures.TestSimpleNode("node-3").SummaryString()),
+				PreemptionDescription: fmt.Sprintf(unknownPreemptionCause, testfixtures.TestSimpleNode("node-3").SummaryString()),
 			},
 		},
 		"unknown cause - gang job": {
 			preemptedJobContext: &context.JobSchedulingContext{
 				JobId:        "job-1",
-				AssignedNode: ittestfixtures.TestSimpleNode("node-3"),
+				AssignedNode: testfixtures.TestSimpleNode("node-3"),
 				Job:          makeJob(t, "job-1", true),
 			},
 			expectedPreemptedJobContext: &context.JobSchedulingContext{
 				JobId:                 "job-1",
-				AssignedNode:          ittestfixtures.TestSimpleNode("node-3"),
+				AssignedNode:          testfixtures.TestSimpleNode("node-3"),
 				Job:                   makeJob(t, "job-1", true),
 				PreemptionDescription: unknownGangPreemptionCause,
 			},
 		},
 		"urgency preemption - single preempting job": {
 			preemptedJobContext: &context.JobSchedulingContext{
 				JobId:        "job-1",
-				AssignedNode: ittestfixtures.TestSimpleNode("node-1"),
+				AssignedNode: testfixtures.TestSimpleNode("node-1"),
 			},
 			expectedPreemptedJobContext: &context.JobSchedulingContext{
 				JobId:                 "job-1",
-				AssignedNode:          ittestfixtures.TestSimpleNode("node-1"),
+				AssignedNode:          testfixtures.TestSimpleNode("node-1"),
 				PreemptionDescription: fmt.Sprintf(urgencyPreemptionTemplate, "job-2"),
 			},
 		},
 		"urgency preemption - multiple preempting jobs": {
 			preemptedJobContext: &context.JobSchedulingContext{
 				JobId:        "job-1",
-				AssignedNode: ittestfixtures.TestSimpleNode("node-2"),
+				AssignedNode: testfixtures.TestSimpleNode("node-2"),
 			},
 			expectedPreemptedJobContext: &context.JobSchedulingContext{
 				JobId:                 "job-1",
-				AssignedNode:          ittestfixtures.TestSimpleNode("node-2"),
+				AssignedNode:          testfixtures.TestSimpleNode("node-2"),
 				PreemptionDescription: fmt.Sprintf(urgencyPreemptionMultiJobTemplate, "job-3,job-4"),
 			},
 		},
 		"fairshare": {
 			preemptedJobContext: &context.JobSchedulingContext{
 				JobId:           "job-1",
-				AssignedNode:    ittestfixtures.TestSimpleNode("node-4"),
+				AssignedNode:    testfixtures.TestSimpleNode("node-4"),
 				PreemptingJobId: "job-7",
 			},
 			expectedPreemptedJobContext: &context.JobSchedulingContext{
 				JobId:                 "job-1",
-				AssignedNode:          ittestfixtures.TestSimpleNode("node-4"),
+				AssignedNode:          testfixtures.TestSimpleNode("node-4"),
 				PreemptingJobId:       "job-7",
 				PreemptionDescription: fmt.Sprintf(fairSharePreemptionTemplate, "job-7"),
 			},