fix: unit test issue (#380)

* fix: preemption unit test

* fix: add write loc for gpu capacity map

* fix: increase timeout for github action run

Author: Code2Life

.vscode/launch.json

@@ -97,6 +97,19 @@
             },
             "program": "${workspaceFolder}/internal/webhook/v1",
             "console": "integratedTerminal"
+        },
+        {
+            "name": "Run Unit Tests - Preemption",
+            "type": "go",
+            "request": "launch",
+            "mode": "test",
+            "env": {
+                "DEBUG_MODE": "true",
+                "GO_TESTING": "true",
+                "ENVTEST_K8S_VERSION": "1.34.0"
+            },
+            "program": "${workspaceFolder}/test/sched",
+            "console": "integratedTerminal"
         }
     ]
 }

internal/controller/suite_test.go

@@ -79,7 +79,7 @@ func TestControllers(t *testing.T) {
 	SetDefaultConsistentlyDuration(5 * time.Second)
 	SetDefaultConsistentlyPollingInterval(250 * time.Millisecond)
 	suiteConfig, reporterConfig := GinkgoConfiguration()
-	suiteConfig.Timeout = 5 * time.Minute
+	suiteConfig.Timeout = 12 * time.Minute
 	RunSpecs(t, "Controller Suite", suiteConfig, reporterConfig)
 }
 

internal/gpuallocator/gpuallocator.go

@@ -40,6 +40,7 @@ import (
 const MaxGPUCounterPerAllocation = 128
 const CleanUpCheckInterval = 3 * time.Minute
 
+var mu sync.Mutex
 var GPUCapacityMap = map[string]tfv1.Resource{}
 
 type Strategy interface {
@@ -921,9 +922,11 @@ func (s *GpuAllocator) handleGPUCreate(ctx context.Context, gpu *tfv1.GPU) {
 
 	if s.gpuStore[key] != nil {
 		if gpu.Status.GPUModel != "" {
+			mu.Lock()
 			if _, exists := GPUCapacityMap[gpu.Status.GPUModel]; !exists {
 				GPUCapacityMap[gpu.Status.GPUModel] = *gpu.Status.Capacity
 			}
+			mu.Unlock()
 		}
 		syncGPUMetadataAndStatusFromCluster(s.gpuStore[key], gpu)
 		log.V(6).Info("GPU already exists in store", "name", key.Name)
@@ -1019,7 +1022,9 @@ func (s *GpuAllocator) addOrUpdateGPUMaps(gpuInMem *tfv1.GPU) {
 	}
 
 	if gpuInMem.Status.GPUModel != "" {
+		mu.Lock()
 		GPUCapacityMap[gpuInMem.Status.GPUModel] = *gpuInMem.Status.Capacity
+		mu.Unlock()
 	}
 }
 

internal/scheduler/expander/handler_test.go

@@ -83,7 +83,7 @@ func (suite *NodeExpanderTestSuite) TearDownSuite() {
 
 func TestNodeExpander(t *testing.T) {
 	suiteConfig, reporterConfig := GinkgoConfiguration()
-	suiteConfig.Timeout = 2 * time.Minute
+	suiteConfig.Timeout = 3 * time.Minute
 	RegisterFailHandler(Fail)
 	if os.Getenv("DEBUG_MODE") == constants.TrueStringValue {
 		SetDefaultEventuallyTimeout(10 * time.Minute)

scripts/benchmark.sh

@@ -1,8 +1,8 @@
 # Code level benchmark 
 # go test -bench=BenchmarkPodWebhookQPS -benchmem -v ./internal/webhook/v1/
 
-# go test -bench=BenchmarkGPUFitPlugin -benchmem ./test/sched/ --benchtime=3s
-# go test -bench=BenchmarkScheduler    -benchmem ./test/sched/ --benchtime=5s
+# go test -bench=BenchmarkGPUFitPlugin -benchmem ./test/sched/ --benchtime=3s -tags=nobench
+# go test -bench=BenchmarkScheduler    -benchmem ./test/sched/ --benchtime=5s -tags=nobench
 
 # Real world benchmark for Mutating Webhook
 cat > /tmp/webhook-body.json << 'EOF'

test/sched/preemption_test.go

@@ -1,9 +1,10 @@
+//go:build !nobench
+
 package sched
 
 import (
 	"context"
 	"fmt"
-	"os"
 	"strings"
 	"testing"
 	"time"
@@ -25,17 +26,6 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/envtest"
 )
 
-// isRunningInGitHubActions checks if the tests are running in GitHub Actions environment
-func isRunningInGitHubActions() bool {
-	if os.Getenv("GITHUB_ACTIONS") == "true" {
-		return true
-	}
-	if os.Getenv("CI") == "true" {
-		return true
-	}
-	return false
-}
-
 // PreemptionTestSuite holds common test setup for preemption tests
 type PreemptionTestSuite struct {
 	ctx            context.Context
@@ -95,6 +85,9 @@ func (pts *PreemptionTestSuite) SetupSuite() {
 	Expect(err).To(Succeed())
 	pts.scheduler = scheduler
 	scheduler.SchedulingQueue.Run(klog.FromContext(ctx))
+	if scheduler.APIDispatcher != nil {
+		scheduler.APIDispatcher.Run(klog.FromContext(ctx))
+	}
 
 	// Start scheduler components
 	cc.EventBroadcaster.StartRecordingToSink(ctx.Done())
@@ -105,6 +98,7 @@ func (pts *PreemptionTestSuite) SetupSuite() {
 
 // TearDownSuite cleans up the test environment
 func (pts *PreemptionTestSuite) TearDownSuite() {
+	time.Sleep(300 * time.Millisecond)
 	if pts.cancel != nil {
 		pts.cancel()
 	}
@@ -122,7 +116,7 @@ func (pts *PreemptionTestSuite) TearDownSuite() {
 // TestPreemption tests comprehensive preemption scenarios
 func TestPreemption(t *testing.T) {
 	suiteConfig, reporterConfig := GinkgoConfiguration()
-	suiteConfig.Timeout = 1 * time.Minute
+	suiteConfig.Timeout = 2 * time.Minute
 	RegisterFailHandler(Fail)
 	RunSpecs(t, "Preemption Test Suite", suiteConfig, reporterConfig)
 }
@@ -140,16 +134,10 @@ var _ = Describe("GPU Resource Preemption", func() {
 	})
 
 	It("should preempt lower priority pods for higher priority ones", func() {
-		if isRunningInGitHubActions() {
-			Skip("Skipping preemption test in GitHub Actions environment")
-		}
 		testGPUResourcePreemption(suite)
 	})
 
 	It("should respect eviction protection periods", func() {
-		if isRunningInGitHubActions() {
-			Skip("Skipping eviction protection test in GitHub Actions environment")
-		}
 		testGPUResourceEvictProtection(suite)
 	})
 })
@@ -193,6 +181,8 @@ func testGPUResourcePreemption(suite *PreemptionTestSuite) {
 	defer func() {
 		_ = suite.k8sClient.Delete(suite.ctx, criticalPriorityPod)
 	}()
+	time.Sleep(10 * time.Millisecond)
+	suite.scheduler.SchedulingQueue.Add(klog.FromContext(suite.ctx), criticalPriorityPod)
 	suite.scheduler.ScheduleOne(suite.ctx)
 
 	// Preemption should be triggered and victims deleted, wait informer sync