[YUNIKORN-2927] Update MockScheduler test case with foreign pod resou…

…rce update (#934) Closes: #934 Signed-off-by: Craig Condit <[email protected]>
apache · Nov 8, 2024 · 8fa300f · 8fa300f
1 parent 52bef67
commit 8fa300f
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 3 deletions.
diff --git a/go.mod b/go.mod
@@ -23,7 +23,7 @@ go 1.22.0
 toolchain go1.22.5
 
 require (
-	github.com/apache/yunikorn-core v0.0.0-20241017135039-079a02dbdfa7
+	github.com/apache/yunikorn-core v0.0.0-20241104184802-6ef347b31c1c
 	github.com/apache/yunikorn-scheduler-interface v0.0.0-20241016105739-f0e241aa0146
 	github.com/google/go-cmp v0.6.0
 	github.com/google/uuid v1.6.0

diff --git a/go.sum b/go.sum
@@ -8,8 +8,8 @@ github.com/NYTimes/gziphandler v1.1.1 h1:ZUDjpQae29j0ryrS0u/B8HZfJBtBQHjqw2rQ2cq
 github.com/NYTimes/gziphandler v1.1.1/go.mod h1:n/CVRwUEOgIxrgPvAQhUUr9oeUtvrhMomdKFjzJNB0c=
 github.com/antlr4-go/antlr/v4 v4.13.0 h1:lxCg3LAv+EUK6t1i0y1V6/SLeUi0eKEKdhQAlS8TVTI=
 github.com/antlr4-go/antlr/v4 v4.13.0/go.mod h1:pfChB/xh/Unjila75QW7+VU4TSnWnnk9UTnmpPaOR2g=
-github.com/apache/yunikorn-core v0.0.0-20241017135039-079a02dbdfa7 h1:PY3kIiQYxsNcs42DK+8b7NxfTvMF0Z6eIuK+aJNWl18=
-github.com/apache/yunikorn-core v0.0.0-20241017135039-079a02dbdfa7/go.mod h1:JA8Uee+D+T9v3p+YznGiGM9cLk5tzX+EM+YYr1TdFYo=
+github.com/apache/yunikorn-core v0.0.0-20241104184802-6ef347b31c1c h1:c0+cVnKSAOiJHC6lNUKEl+tt7lZLIEfqv0cPaTI//4U=
+github.com/apache/yunikorn-core v0.0.0-20241104184802-6ef347b31c1c/go.mod h1:JA8Uee+D+T9v3p+YznGiGM9cLk5tzX+EM+YYr1TdFYo=
 github.com/apache/yunikorn-scheduler-interface v0.0.0-20241016105739-f0e241aa0146 h1:CZ4U7y19YSxNJVBNox3DahhuoxDL++naBl/kj+kqVFc=
 github.com/apache/yunikorn-scheduler-interface v0.0.0-20241016105739-f0e241aa0146/go.mod h1:co3uU98sj1CUTPNTM13lTyi+CY0DOgDndDW2KiUjktU=
 github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5 h1:0CwZNZbxp69SHPdPJAN/hZIm0C4OItdklCFmMRWYpio=

diff --git a/pkg/shim/scheduler_mock_test.go b/pkg/shim/scheduler_mock_test.go
@@ -336,6 +336,25 @@ func (fc *MockScheduler) waitAndAssertForeignAllocationInCore(partition, allocat
 	}, time.Second, 5*time.Second)
 }
 
+func (fc *MockScheduler) waitAndAssertForeignAllocationResources(partition, allocationID, nodeID string, expected *si.Resource) error {
+	return utils.WaitForCondition(func() bool {
+		node := fc.coreContext.Scheduler.GetClusterContext().GetNode(nodeID, partition)
+		if node == nil {
+			log.Log(log.Test).Warn("Node not found", zap.String("node ID", nodeID))
+			return false
+		}
+		allocs := node.GetForeignAllocations()
+		for _, alloc := range allocs {
+			if alloc.GetAllocationKey() == allocationID {
+				current := alloc.GetAllocatedResource().ToProto()
+				return common.Equals(expected, current)
+			}
+		}
+
+		return false
+	}, time.Second, 5*time.Second)
+}
+
 func (fc *MockScheduler) getApplicationFromCore(appID, partition string) *objects.Application {
 	return fc.coreContext.Scheduler.GetClusterContext().GetApplication(appID, partition)
 }

diff --git a/pkg/shim/scheduler_test.go b/pkg/shim/scheduler_test.go
@@ -350,6 +350,21 @@ partitions:
 	err = cluster.waitAndAssertForeignAllocationInCore(partitionName, "foreign-2", "node-1", true)
 	assert.NilError(t, err)
 
+	// update pod resources
+	pod1Copy := pod1.DeepCopy()
+	pod1Copy.Spec.Containers[0].Resources.Requests[siCommon.Memory] = *resource.NewQuantity(500, resource.DecimalSI)
+	pod1Copy.Spec.Containers[0].Resources.Requests[siCommon.CPU] = *resource.NewMilliQuantity(2000, resource.DecimalSI)
+
+	cluster.UpdatePod(pod1, pod1Copy)
+	expectedUsage := common.NewResourceBuilder().
+		AddResource(siCommon.Memory, 500).
+		AddResource(siCommon.CPU, 2).
+		AddResource("pods", 1).
+		Build()
+	err = cluster.waitAndAssertForeignAllocationResources(partitionName, "foreign-1", "node-1", expectedUsage)
+	assert.NilError(t, err)
+
+	// delete pods
 	cluster.DeletePod(pod1)
 	cluster.DeletePod(pod2)