fix: unit test

Author: Code2Life

cmd/main.go

@@ -235,9 +235,6 @@ func main() {
 
 	metricsRecorder := startMetricsRecorder(enableLeaderElection, mgr, gpuPricingMap)
 
-	// Initialize GPU allocator and set up watches
-	allocator, portAllocator := startTensorFusionAllocators(ctx, mgr)
-
 	// Initialize Index allocator for Device Plugin communication
 	indexAllocator, err := indexallocator.NewIndexAllocator(ctx, mgr.GetClient())
 	if err != nil {
@@ -246,6 +243,9 @@ func main() {
 	}
 	_ = indexAllocator.SetupWithManager(ctx, mgr)
 
+	// Initialize GPU allocator and set up watches
+	allocator, portAllocator := startTensorFusionAllocators(ctx, mgr, indexAllocator)
+
 	ensureLeaderInfoConfigMap(mgr)
 
 	startAutoScaler(mgr, allocator)
@@ -286,8 +286,9 @@ func addHealthCheckAPI(mgr manager.Manager) {
 func startTensorFusionAllocators(
 	ctx context.Context,
 	mgr manager.Manager,
+	indexAllocator *indexallocator.IndexAllocator,
 ) (*gpuallocator.GpuAllocator, *portallocator.PortAllocator) {
-	allocator := gpuallocator.NewGpuAllocator(ctx, mgr.GetClient(), 10*time.Second)
+	allocator := gpuallocator.NewGpuAllocator(ctx, indexAllocator, mgr.GetClient(), 10*time.Second)
 	if err := allocator.SetupWithManager(ctx, mgr); err != nil {
 		setupLog.Error(err, "unable to set up GPU allocator watches")
 		os.Exit(1)

internal/autoscaler/autoscaler_suite_test.go

@@ -155,7 +155,7 @@ var _ = BeforeSuite(func() {
 		WorkerUnitPriceMap: make(map[string]map[string]metrics.RawBillingPricing),
 	}
 
-	allocator = gpuallocator.NewGpuAllocator(ctx, mgr.GetClient(), 150*time.Millisecond)
+	allocator = gpuallocator.NewGpuAllocator(ctx, nil, mgr.GetClient(), 150*time.Millisecond)
 	err = allocator.SetupWithManager(ctx, mgr)
 	Expect(err).ToNot(HaveOccurred())
 

internal/controller/suite_test.go

@@ -156,7 +156,7 @@ var _ = BeforeSuite(func() {
 		WorkerUnitPriceMap: make(map[string]map[string]metrics.RawBillingPricing),
 	}
 
-	allocator = gpuallocator.NewGpuAllocator(ctx, mgr.GetClient(), 150*time.Millisecond)
+	allocator = gpuallocator.NewGpuAllocator(ctx, nil, mgr.GetClient(), 150*time.Millisecond)
 	err = allocator.SetupWithManager(ctx, mgr)
 	Expect(err).ToNot(HaveOccurred())
 

internal/gpuallocator/gpuallocator.go

@@ -150,7 +150,12 @@ type GpuAllocator struct {
 	indexAllocator *indexallocator.IndexAllocator
 }
 
-func NewGpuAllocator(ctx context.Context, client client.Client, syncInterval time.Duration) *GpuAllocator {
+func NewGpuAllocator(
+	ctx context.Context,
+	indexAllocator *indexallocator.IndexAllocator,
+	client client.Client,
+	syncInterval time.Duration,
+) *GpuAllocator {
 	log := log.FromContext(ctx)
 
 	if client == nil {
@@ -166,6 +171,15 @@ func NewGpuAllocator(ctx context.Context, client client.Client, syncInterval tim
 	// Create quota store
 	quotaStore := quota.NewQuotaStore(client, ctx)
 
+	if indexAllocator == nil {
+		newIndexAllocator, err := indexallocator.NewIndexAllocator(ctx, client)
+		if err != nil {
+			log.Error(err, "Failed to create index allocator")
+			return nil
+		}
+		indexAllocator = newIndexAllocator
+	}
+
 	allocator := &GpuAllocator{
 		Client:          client,
 		filterRegistry:  baseRegistry,
@@ -178,6 +192,7 @@ func NewGpuAllocator(ctx context.Context, client client.Client, syncInterval tim
 		dirtyQueue:      make(map[types.NamespacedName]struct{}),
 		ctx:             ctx,
 
+		indexAllocator:         indexAllocator,
 		uniqueAllocation:       make(map[string]*tfv1.AllocRequest, 512),
 		uniqueDeallocation:     make(map[string]struct{}, 512),
 		podNamespaceNsToPodUID: make(map[string]string, 512),

internal/gpuallocator/gpuallocator_test.go

@@ -66,7 +66,7 @@ var _ = Describe("GPU Allocator", func() {
 	}
 
 	BeforeEach(func() {
-		allocator = NewGpuAllocator(ctx, k8sClient, 150*time.Millisecond)
+		allocator = NewGpuAllocator(ctx, nil, k8sClient, 150*time.Millisecond)
 		err := allocator.SetupWithManager(ctx, mgr)
 		Expect(err).NotTo(HaveOccurred())
 		<-allocator.initializedCh

internal/gpuallocator/quota_consolidated_test.go

@@ -377,7 +377,7 @@ var _ = Describe("GPUAllocator Quota Integration", func() {
 			Build()
 
 		ctx := context.Background()
-		allocator := NewGpuAllocator(ctx, client, 0)
+		allocator := NewGpuAllocator(ctx, nil, client, 0)
 
 		initAllocator(allocator)
 
@@ -415,7 +415,7 @@ var _ = Describe("GPUAllocator Quota Integration", func() {
 			Build()
 
 		ctx := context.Background()
-		allocator := NewGpuAllocator(ctx, client, 0)
+		allocator := NewGpuAllocator(ctx, nil, client, 0)
 
 		initAllocator(allocator)
 
@@ -451,7 +451,7 @@ var _ = Describe("GPUAllocator Concurrent Quota Enforcement", func() {
 			Build()
 
 		ctx := context.Background()
-		allocator := NewGpuAllocator(ctx, client, 0)
+		allocator := NewGpuAllocator(ctx, nil, client, 0)
 
 		initAllocator(allocator)
 
@@ -539,7 +539,7 @@ var _ = Describe("GPUAllocator Quota Reconciliation", func() {
 			Build()
 
 		ctx := context.Background()
-		allocator := NewGpuAllocator(ctx, client, 0)
+		allocator := NewGpuAllocator(ctx, nil, client, 0)
 
 		initAllocator(allocator)
 

internal/scheduler/expander/handler_test.go

@@ -56,7 +56,7 @@ func (suite *NodeExpanderTestSuite) SetupSuite() {
 	Expect(suite.k8sClient.Create(ctx, ns)).To(Succeed())
 
 	// Setup proper allocator for testing
-	suite.allocator = gpuallocator.NewGpuAllocator(ctx, suite.k8sClient, time.Second)
+	suite.allocator = gpuallocator.NewGpuAllocator(ctx, nil, suite.k8sClient, time.Second)
 	err := suite.allocator.InitGPUAndQuotaStore()
 	if err != nil {
 		// For test environments, we can ignore some initialization errors

internal/scheduler/gpuresources/gpuresources_test.go

@@ -259,7 +259,7 @@ func (s *GPUResourcesSuite) SetupTest() {
 	s.NoError(err)
 	s.fwk = fwk
 
-	s.allocator = gpuallocator.NewGpuAllocator(s.ctx, s.client, time.Second)
+	s.allocator = gpuallocator.NewGpuAllocator(s.ctx, nil, s.client, time.Second)
 	err = s.allocator.InitGPUAndQuotaStore()
 	s.NoError(err)
 	s.allocator.ReconcileAllocationState()

internal/webhook/v1/pod_webhook.go

@@ -310,7 +310,7 @@ func (m *TensorFusionPodMutator) createOrUpdateWorkload(
 }
 
 func (m *TensorFusionPodMutator) patchTFClient(
-	_ctx context.Context,
+	ctx context.Context,
 	pod *corev1.Pod,
 	pool *tfv1.GPUPool,
 	isLocalGPU bool,

test/sched/setup.go

@@ -387,7 +387,7 @@ func setupFrameworkAndPlugin(
 func setupAllocator(
 	b *testing.B, ctx context.Context, client client.Client,
 ) *gpuallocator.GpuAllocator {
-	allocator := gpuallocator.NewGpuAllocator(ctx, client, time.Second)
+	allocator := gpuallocator.NewGpuAllocator(ctx, nil, client, time.Second)
 	require.NoError(b, allocator.InitGPUAndQuotaStore())
 	allocator.ReconcileAllocationState()
 	allocator.SetAllocatorReady()