tidy

Signed-off-by: Robert Smith <[email protected]>
armadaproject · Dec 30, 2024 · 0cf3067 · 0cf3067
1 parent bfc10f3
commit 0cf3067
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 33 deletions.
diff --git a/internal/scheduler/internaltypes/node_factory.go b/internal/scheduler/internaltypes/node_factory.go
@@ -122,8 +122,8 @@ func (f *NodeFactory) AddLabels(nodes []*Node, extraLabels map[string]string) []
 			newLabels,
 			f.indexedTaints,
 			f.indexedNodeLabels,
-			node.totalResources,
-			node.unallocatableResources,
+			node.GetTotalResources(),
+			node.GetUnallocatableResources(),
 			node.AllocatableByPriority,
 		)
 	}
@@ -143,8 +143,8 @@ func (f *NodeFactory) AddTaints(nodes []*Node, extraTaints []v1.Taint) []*Node {
 			node.GetLabels(),
 			f.indexedTaints,
 			f.indexedNodeLabels,
-			node.totalResources,
-			node.unallocatableResources,
+			node.GetTotalResources(),
+			node.GetUnallocatableResources(),
 			node.AllocatableByPriority,
 		)
 	}

diff --git a/internal/scheduler/internaltypes/resource_list_map_util_test.go b/internal/scheduler/internaltypes/resource_list_map_util_test.go
@@ -88,6 +88,46 @@ func TestNewAllocatableByPriorityAndResourceType(t *testing.T) {
 	assert.Equal(t, int64(2000), result[2].GetByNameZeroIfMissing("cpu"))
 }
 
+func TestMarkAllocated(t *testing.T) {
+	factory := testFactory()
+	m := map[int32]ResourceList{
+		1: testResourceList(factory, "10", "10Gi"),
+		2: testResourceList(factory, "20", "20Gi"),
+		3: testResourceList(factory, "30", "30Gi"),
+		4: testResourceList(factory, "40", "40Gi"),
+	}
+
+	expected := map[int32]ResourceList{
+		1: testResourceList(factory, "8", "8Gi"),
+		2: testResourceList(factory, "18", "18Gi"),
+		3: testResourceList(factory, "30", "30Gi"),
+		4: testResourceList(factory, "40", "40Gi"),
+	}
+
+	MarkAllocated(m, 2, testResourceList(factory, "2", "2Gi"))
+	assert.Equal(t, expected, m)
+}
+
+func TestMarkAllocatable(t *testing.T) {
+	factory := testFactory()
+	m := map[int32]ResourceList{
+		1: testResourceList(factory, "10", "10Gi"),
+		2: testResourceList(factory, "20", "20Gi"),
+		3: testResourceList(factory, "30", "30Gi"),
+		4: testResourceList(factory, "40", "40Gi"),
+	}
+
+	expected := map[int32]ResourceList{
+		1: testResourceList(factory, "12", "12Gi"),
+		2: testResourceList(factory, "22", "22Gi"),
+		3: testResourceList(factory, "30", "30Gi"),
+		4: testResourceList(factory, "40", "40Gi"),
+	}
+
+	MarkAllocatable(m, 2, testResourceList(factory, "2", "2Gi"))
+	assert.Equal(t, expected, m)
+}
+
 func testMapAllPositive(factory *ResourceListFactory) map[string]ResourceList {
 	return map[string]ResourceList{
 		"a": testResourceList(factory, "1", "1Ki"),

diff --git a/internal/scheduler/nodedb/nodedb_test.go b/internal/scheduler/nodedb/nodedb_test.go
@@ -358,12 +358,13 @@ func TestScheduleIndividually(t *testing.T) {
 			ExpectSuccess: []bool{false, false, true},
 		},
 		"node selector": {
-			Nodes: testfixtures.TestNodeFactory.AddLabels(
-				testfixtures.ItN32CpuNodes(1, testfixtures.TestPriorities),
-				map[string]string{
-					"key": "value",
-				},
-			),
+			Nodes: append(testfixtures.ItN32CpuNodes(1, testfixtures.TestPriorities),
+				testfixtures.TestNodeFactory.AddLabels(
+					testfixtures.ItN32CpuNodes(1, testfixtures.TestPriorities),
+					map[string]string{
+						"key": "value",
+					},
+				)...),
 			Jobs: testfixtures.WithNodeSelectorJobs(
 				map[string]string{
 					"key": "value",
@@ -898,28 +899,6 @@ func newNodeDbWithNodes(nodes []*internaltypes.Node) (*NodeDb, error) {
 	return nodeDb, nil
 }
 
-func itNewNodeDbWithNodes(nodes []*internaltypes.Node) (*NodeDb, error) {
-	nodeDb, err := NewNodeDb(
-		testfixtures.TestPriorityClasses,
-		testfixtures.TestResources,
-		testfixtures.TestIndexedTaints,
-		testfixtures.TestIndexedNodeLabels,
-		testfixtures.TestWellKnownNodeTypes,
-		testfixtures.TestResourceListFactory,
-	)
-	if err != nil {
-		return nil, err
-	}
-	txn := nodeDb.Txn(true)
-	for _, node := range nodes {
-		if err = nodeDb.CreateAndInsertWithJobDbJobsWithTxn(txn, nil, node); err != nil {
-			return nil, err
-		}
-	}
-	txn.Commit()
-	return nodeDb, nil
-}
-
 func BenchmarkNodeDbStringFromPodRequirementsNotMetReason(b *testing.B) {
 	nodeDb := &NodeDb{
 		podRequirementsNotMetReasonStringCache: make(map[uint64]string, 128),

diff --git a/internal/scheduler/nodedb/nodeiteration_test.go b/internal/scheduler/nodedb/nodeiteration_test.go
@@ -38,7 +38,7 @@ func TestNodesIterator(t *testing.T) {
 			for i, node := range tc.Nodes {
 				indexById[node.GetId()] = i
 			}
-			nodeDb, err := itNewNodeDbWithNodes(tc.Nodes)
+			nodeDb, err := newNodeDbWithNodes(tc.Nodes)
 			if !assert.NoError(t, err) {
 				return
 			}