fix: unit test issues (#302)

Author: Code2Life

cmd/nodediscovery/main.go

@@ -179,14 +179,16 @@ func main() {
 }
 
 // Use proper patch-based update with retry on conflict
-func patchGPUNodeStatus(k8sClient client.Client, ctx context.Context, gpunode *tfv1.GPUNode, totalTFlops resource.Quantity, totalVRAM resource.Quantity, count int32, allDeviceIDs []string) error {
+func patchGPUNodeStatus(k8sClient client.Client, ctx context.Context,
+	gpunode *tfv1.GPUNode, totalTFlops resource.Quantity, totalVRAM resource.Quantity,
+	count int32, allDeviceIDs []string) error {
 
 	currentGPUNode := &tfv1.GPUNode{}
 	if err := k8sClient.Get(ctx, client.ObjectKeyFromObject(gpunode), currentGPUNode); err != nil {
 		return err
 	}
 	patch := client.MergeFrom(currentGPUNode.DeepCopy())
-	updateGPUNodeStatus(&currentGPUNode.Status, totalTFlops, totalVRAM, int32(count), allDeviceIDs)
+	updateGPUNodeStatus(&currentGPUNode.Status, totalTFlops, totalVRAM, count, allDeviceIDs)
 	return k8sClient.Status().Patch(ctx, currentGPUNode, patch)
 }
 

cmd/nodediscovery/main_test.go

@@ -181,7 +181,8 @@ func TestPatchGPUNodeStatus(t *testing.T) {
 				assert.Equal(t, resource.MustParse("24Gi"), patchedNode.Status.TotalVRAM)
 				assert.Equal(t, int32(6), patchedNode.Status.TotalGPUs)
 				assert.Equal(t, int32(6), patchedNode.Status.ManagedGPUs)
-				assert.Equal(t, []string{"gpu-0", "gpu-1", "gpu-2", "gpu-3", "gpu-4", "gpu-5"}, patchedNode.Status.ManagedGPUDeviceIDs)
+				assert.Equal(t, []string{"gpu-0", "gpu-1", "gpu-2", "gpu-3", "gpu-4", "gpu-5"},
+					patchedNode.Status.ManagedGPUDeviceIDs)
 				// Verify existing phase was preserved
 				assert.Equal(t, tfv1.TensorFusionGPUNodePhaseRunning, patchedNode.Status.Phase)
 			},
@@ -368,7 +369,8 @@ func TestPatchGPUNodeStatus_Integration(t *testing.T) {
 	for i, update := range updates {
 		t.Run(fmt.Sprintf("update_%d", i+1), func(t *testing.T) {
 			// Apply the patch
-			err := patchGPUNodeStatus(k8sClient, ctx, gpuNode, update.totalTFlops, update.totalVRAM, update.count, update.allDeviceIDs)
+			err := patchGPUNodeStatus(k8sClient, ctx, gpuNode, update.totalTFlops,
+				update.totalVRAM, update.count, update.allDeviceIDs)
 			assert.NoError(t, err, "Patch should succeed")
 
 			// Verify the update was applied

internal/controller/gpupool_node_provision.go

@@ -145,23 +145,23 @@ func (r *GPUPoolReconciler) reconcilePoolCapacityWithProvisioner(ctx context.Con
 
 	// lock the pool before next node scaling up loop, add assumed scaling resources util all pending nodeClaims are running
 	newCreatedNodes := map[string]tfv1.Resource{}
-	for _, node := range gpuNodeParams {
-		go func(node tfv1.GPUNodeClaimSpec) {
+	for _, nodeClaim := range gpuNodeParams {
+		go func(nodeClaim tfv1.GPUNodeClaimSpec) {
 			defer wg.Done()
 
 			// Create GPUNode custom resource immediately and GPUNode controller will watch the K8S node to be ready
 			// Persist the status to GPUNode to avoid duplicated creation in next reconciliation
 			// If the K8S node never be ready after some time, the GPUNode will be deleted, then the Pool reconcile loop can scale up and meet the capacity constraint again
 
-			costPerHour, pricingErr := provider.GetInstancePricing(node.InstanceType, node.CapacityType, node.Region)
+			costPerHour, pricingErr := provider.GetInstancePricing(nodeClaim.InstanceType, nodeClaim.CapacityType, nodeClaim.Region)
 			if pricingErr != nil {
 				errList = append(errList, pricingErr)
 				return
 			}
 
 			gpuNodeClaimRes := &tfv1.GPUNodeClaim{
 				ObjectMeta: metav1.ObjectMeta{
-					Name: node.NodeName,
+					Name: nodeClaim.NodeName,
 					Labels: map[string]string{
 						constants.LabelKeyOwner:        pool.Name,
 						constants.LabelKeyClusterOwner: cluster.Name,
@@ -174,19 +174,20 @@ func (r *GPUPoolReconciler) reconcilePoolCapacityWithProvisioner(ctx context.Con
 						constants.PricingAnnotation: strconv.FormatFloat(costPerHour, 'f', 6, 64),
 					},
 				},
-				Spec: node,
+				Spec: nodeClaim,
 			}
 			_ = controllerutil.SetControllerReference(pool, gpuNodeClaimRes, r.Scheme)
 			err := r.Create(ctx, gpuNodeClaimRes)
 			if err != nil {
 				errList = append(errList, err)
 				return
 			}
-			newCreatedNodes[node.NodeName] = tfv1.Resource{
-				Tflops: node.TFlopsOffered,
-				Vram:   node.VRAMOffered,
+			log.Info("Created new GPUNode claim", "gpuNodeClaimName", nodeClaim.NodeName)
+			newCreatedNodes[nodeClaim.NodeName] = tfv1.Resource{
+				Tflops: nodeClaim.TFlopsOffered,
+				Vram:   nodeClaim.VRAMOffered,
 			}
-		}(node)
+		}(nodeClaim)
 	}
 
 	wg.Wait()

internal/controller/suite_test.go

@@ -506,24 +506,22 @@ func (c *TensorFusionEnv) AddMockGPU4ProvisionedNodes(gpuNodeClaimList *tfv1.GPU
 				if err := k8sClient.Status().Update(ctx, latest); err != nil {
 					return err
 				}
-
-				// update GPUNode status to trigger node level reconcile, simulate node discovery job
-				if gpuNode.Status.Phase == "" || gpuNode.Status.TotalGPUs == 0 {
-					gpuNode.Status = tfv1.GPUNodeStatus{
-						Phase:       tfv1.TensorFusionGPUNodePhasePending,
-						TotalGPUs:   1,
-						ManagedGPUs: 1,
-						TotalTFlops: gpuNodeClaim.Spec.TFlopsOffered,
-						TotalVRAM:   gpuNodeClaim.Spec.VRAMOffered,
-					}
-					if err := k8sClient.Status().Update(ctx, gpuNode); err != nil {
-						return err
-					}
-				}
 				return nil
 			})
 			Expect(err).Should(Succeed())
 		}
+
+		// update GPUNode status to trigger node level reconcile, simulate node discovery job
+		if gpuNode.Status.Phase == "" || gpuNode.Status.TotalGPUs == 0 {
+			gpuNode.Status = tfv1.GPUNodeStatus{
+				Phase:       tfv1.TensorFusionGPUNodePhasePending,
+				TotalGPUs:   1,
+				ManagedGPUs: 1,
+				TotalTFlops: gpuNodeClaim.Spec.TFlopsOffered,
+				TotalVRAM:   gpuNodeClaim.Spec.VRAMOffered,
+			}
+			Expect(k8sClient.Status().Update(ctx, gpuNode)).Should(Succeed())
+		}
 	}
 }