From 7c51e82f09720b014159e0ba2f48d5e600a1f96c Mon Sep 17 00:00:00 2001
From: SP12893678 <36910625+SP12893678@users.noreply.github.com>
Date: Mon, 9 Sep 2024 14:22:52 -0500
Subject: [PATCH] [YUNIKORN-2862] Improve nodesorting test coverage (#967)

Closes: #967

Signed-off-by: Craig Condit <ccondit@apache.org>
---
 pkg/scheduler/objects/nodesorting.go      |  4 --
 pkg/scheduler/objects/nodesorting_test.go | 46 +++++++++++++++++++++--
 2 files changed, 43 insertions(+), 7 deletions(-)

diff --git a/pkg/scheduler/objects/nodesorting.go b/pkg/scheduler/objects/nodesorting.go
index e8674b876..7ac6801a7 100644
--- a/pkg/scheduler/objects/nodesorting.go
+++ b/pkg/scheduler/objects/nodesorting.go
@@ -131,10 +131,6 @@ func NewNodeSortingPolicy(policyType string, resourceWeights map[string]float64)
 		sp = fairnessNodeSortingPolicy{
 			resourceWeights: weights,
 		}
-	default:
-		sp = fairnessNodeSortingPolicy{
-			resourceWeights: weights,
-		}
 	}
 
 	log.Log(log.SchedNode).Debug("new node sorting policy added",
diff --git a/pkg/scheduler/objects/nodesorting_test.go b/pkg/scheduler/objects/nodesorting_test.go
index 35ba9063e..896b0e709 100644
--- a/pkg/scheduler/objects/nodesorting_test.go
+++ b/pkg/scheduler/objects/nodesorting_test.go
@@ -75,10 +75,20 @@ partitions:
 	if !ok {
 		t.Fatal("Didn't get expected policy")
 	}
+	resourceWeights := policy.ResourceWeights()
+	assert.Equal(t, 2, len(resourceWeights), "Wrong size of resourceWeights")
+	assert.Equal(t, resourceWeights["vcore"], 1.0, "Wrong weight for vcore")
+	assert.Equal(t, resourceWeights["memory"], 1.0, "Wrong weight for memory")
 
-	assert.Equal(t, 2, len(policy.resourceWeights), "Wrong size of resourceWeights")
-	assert.Equal(t, policy.resourceWeights["vcore"], 1.0, "Wrong weight for vcore")
-	assert.Equal(t, policy.resourceWeights["memory"], 1.0, "Wrong weight for memory")
+	// case: bin packing policy
+	binPackingPolicy, ok := NewNodeSortingPolicy("binpacking", weights).(binPackingNodeSortingPolicy)
+	if !ok {
+		t.Fatal("Didn't get expected policy")
+	}
+	resourceWeights = binPackingPolicy.ResourceWeights()
+	assert.Equal(t, 2, len(resourceWeights), "Wrong size of resourceWeights")
+	assert.Equal(t, resourceWeights["vcore"], 1.0, "Wrong weight for vcore")
+	assert.Equal(t, resourceWeights["memory"], 1.0, "Wrong weight for memory")
 }
 
 func TestInvalidResourceWeightsFromConfig(t *testing.T) {
@@ -287,3 +297,33 @@ func TestSortPolicy(t *testing.T) {
 	assert.Equal(t, node1.NodeID, nodes[0].NodeID, "wrong initial node (fair, node2 half-filled)")
 	assert.Equal(t, node2.NodeID, nodes[1].NodeID, "wrong second node (binpacking, node2 half-filled")
 }
+
+func TestAbsResourceUsage(t *testing.T) {
+	// case: zero weight & NaN shares
+	nc := NewNodeCollection("test")
+	weights := map[string]float64{
+		"vcore":  4.0,
+		"memory": 0,
+	}
+	fair, ok := NewNodeSortingPolicy("fair", weights).(fairnessNodeSortingPolicy)
+	if !ok {
+		t.Fatal("Didn't get fair policy")
+	}
+
+	nc.SetNodeSortingPolicy(fair)
+	totalRes := resources.NewResourceFromMap(map[string]resources.Quantity{"vcore": 0, "memory": 16000})
+
+	proto1 := newProto("test1", totalRes, nil, map[string]string{})
+	node1 := NewNode(proto1)
+	if err := nc.AddNode(node1); err != nil {
+		t.Fatal("Failed to add node1")
+	}
+
+	// add allocations
+	res1 := resources.NewResourceFromMap(map[string]resources.Quantity{"vcore": 0, "memory": 12000})
+	alloc1 := newAllocation("test-app-1", "test1", res1)
+	node1.AddAllocation(alloc1)
+
+	// node1 w/ fair: 400% vcore, 0% memory => ((0 * 4) + (.75 * NaN)) / 0 = 0
+	assert.Equal(t, 0.0, fair.ScoreNode(node1), "Wrong fair score for node1")
+}