fix: unit test bug, init/reconcile gpu store order issue (#231)

* fix: unit test bug, init/reconcile gpu store order issue

* fix: mock db won't require close() function

Author: Code2Life

internal/alert/evaluator_test.go

@@ -16,7 +16,7 @@ import (
 )
 
 // Mock data for testing
-func setupMockDB(t *testing.T) (*sql.DB, sqlmock.Sqlmock, *metrics.TimeSeriesDB) {
+func setupMockDB(t *testing.T) (sqlmock.Sqlmock, *metrics.TimeSeriesDB) {
 	db, mock, err := sqlmock.New()
 	require.NoError(t, err)
 
@@ -27,7 +27,7 @@ func setupMockDB(t *testing.T) (*sql.DB, sqlmock.Sqlmock, *metrics.TimeSeriesDB)
 	require.NoError(t, err)
 
 	tsdb := &metrics.TimeSeriesDB{DB: gormDB}
-	return db, mock, tsdb
+	return mock, tsdb
 }
 
 func createTestRule(name string) *Rule {
@@ -48,12 +48,7 @@ func createTestRule(name string) *Rule {
 }
 
 func TestAlertEvaluator(t *testing.T) {
-	db, mock, tsdb := setupMockDB(t)
-	defer func() {
-		if err := db.Close(); err != nil {
-			t.Error("failed to close db", err)
-		}
-	}()
+	mock, tsdb := setupMockDB(t)
 
 	evaluator := newAlertEvaluator(tsdb)
 
@@ -155,13 +150,7 @@ func TestAlertEvaluator(t *testing.T) {
 	})
 
 	t.Run("process query results empty results", func(t *testing.T) {
-		db, mock, tsdb := setupMockDB(t)
-		defer func() {
-			if err := db.Close(); err != nil {
-				t.Error("failed to close db", err)
-			}
-		}()
-
+		mock, tsdb := setupMockDB(t)
 		rule := createTestRule("test-rule")
 		evaluator := newAlertEvaluator(tsdb)
 
@@ -218,20 +207,14 @@ func TestAlertEvaluator(t *testing.T) {
 
 	// Test edge cases
 	t.Run("evaluate database error", func(t *testing.T) {
-		db, mock, tsdb := setupMockDB(t)
-		defer func() {
-			if err := db.Close(); err != nil {
-				t.Error("failed to close db", err)
-			}
-		}()
-
+		mock, tsdb := setupMockDB(t)
 		rule := createTestRule("test-rule")
 
 		evaluator := newAlertEvaluator(tsdb)
 
 		// Mock database error
 		mock.ExpectQuery("SELECT value, instance, job FROM metrics").
-			WillReturnError(sql.ErrConnDone)
+			WillReturnError(sql.ErrNoRows)
 
 		alerts, err := evaluator.evaluate(rule)
 		assert.Error(t, err)

internal/gpuallocator/gpuallocator.go

@@ -327,12 +327,13 @@ func (s *GpuAllocator) SetupWithManager(ctx context.Context, mgr manager.Manager
 			log.Error(err, "Failed to initialize GPU store")
 			return err
 		}
-		readyCh <- struct{}{}
+		close(readyCh)
 		return nil
 	}))
 
 	go func() {
 		<-mgr.Elected()
+		<-readyCh
 		// reconcile allocation state based on existing workers, run only when it's elected as leader
 		// and only if it's leader, it will start allocating resources to workers, and start sync loop here
 		s.reconcileAllocationState(ctx)
@@ -435,12 +436,37 @@ func (s *GpuAllocator) syncToK8s(ctx context.Context) {
 	}
 
 	for nodeName := range dirtyNodes {
-		// Refer https://datatracker.ietf.org/doc/html/rfc6901#section-3 encode `/` as `~1`
-		patch := []byte(`[{
-			"op": "add",
-			"path": "/metadata/annotations/` + strings.ReplaceAll(constants.GPULastReportTimeAnnotationKey, "/", "~1") + `",
-			"value": "` + time.Now().Format(time.RFC3339) + `"
-		}]`)
+		// First, get the current node to check if annotations exist
+		node := &tfv1.GPUNode{}
+		nodeKey := client.ObjectKey{Name: nodeName}
+		if err := s.Get(ctx, nodeKey, node); err != nil {
+			log.Error(err, "Failed to get GPU node for updating last report time", "node", nodeName)
+			continue
+		}
+
+		var patch []byte
+		timeValue := time.Now().Format(time.RFC3339)
+		encodedKey := strings.ReplaceAll(constants.GPULastReportTimeAnnotationKey, "/", "~1")
+
+		// Check if annotations already exist
+		if node.Annotations == nil {
+			// Create annotations if they don't exist
+			patch = []byte(`[{
+				"op": "add",
+				"path": "/metadata/annotations",
+				"value": {
+					"` + constants.GPULastReportTimeAnnotationKey + `": "` + timeValue + `"
+				}
+			}]`)
+		} else {
+			// Add to existing annotations
+			patch = []byte(`[{
+				"op": "add",
+				"path": "/metadata/annotations/` + encodedKey + `",
+				"value": "` + timeValue + `"
+			}]`)
+		}
+
 		err := retry.RetryOnConflict(retry.DefaultBackoff, func() error {
 			return s.Patch(ctx, &tfv1.GPUNode{
 				ObjectMeta: metav1.ObjectMeta{
@@ -449,7 +475,7 @@ func (s *GpuAllocator) syncToK8s(ctx context.Context) {
 			}, client.RawPatch(types.JSONPatchType, patch))
 		})
 		if err != nil {
-			log.Error(err, "Failed to update GPU node last report time, will retry later", "node", nodeName)
+			log.Error(err, "Failed to update GPU node last report time, allocation state may be inconsistent", "node", nodeName)
 		}
 	}
 }