Apply observedGeneration fix to Workspace and InferenceSet controllers

- Add observedGeneration check to UpdateStatusConditionIfNotMatch in workspace.go
- Add observedGeneration check to UpdateStatusConditionIfNotMatch in inferenceset.go
- Update workspace_test.go with observedGeneration in matching condition
- Update inferenceset_test.go with observedGeneration in matching condition
- Add test cases for stale observedGeneration scenario in both test files

Fixes the same kubectl wait timeout issue for Workspace and InferenceSet resources.

Co-authored-by: andyzhangx <[email protected]>

Author: Copilot

pkg/utils/inferenceset/inferenceset.go

@@ -37,7 +37,7 @@ import (
 func UpdateStatusConditionIfNotMatch(ctx context.Context, c client.Client, iObj *kaitov1alpha1.InferenceSet, cType kaitov1alpha1.ConditionType,
 	cStatus metav1.ConditionStatus, cReason, cMessage string) error {
 	if curCondition := meta.FindStatusCondition(iObj.Status.Conditions, string(cType)); curCondition != nil {
-		if curCondition.Status == cStatus && curCondition.Reason == cReason && curCondition.Message == cMessage {
+		if curCondition.Status == cStatus && curCondition.Reason == cReason && curCondition.Message == cMessage && curCondition.ObservedGeneration == iObj.GetGeneration() {
 			// Nothing to change
 			return nil
 		}

pkg/utils/inferenceset/inferenceset_test.go

@@ -46,10 +46,11 @@ func TestUpdateStatusConditionIfNotMatch(t *testing.T) {
 			Status: kaitov1alpha1.InferenceSetStatus{
 				Conditions: []metav1.Condition{
 					{
-						Type:    string(kaitov1beta1.ConditionTypeResourceStatus),
-						Status:  metav1.ConditionTrue,
-						Reason:  "ResourcesReady",
-						Message: "All resources are ready",
+						Type:               string(kaitov1beta1.ConditionTypeResourceStatus),
+						Status:             metav1.ConditionTrue,
+						Reason:             "ResourcesReady",
+						Message:            "All resources are ready",
+						ObservedGeneration: 1,
 					},
 				},
 			},
@@ -275,6 +276,55 @@ func TestUpdateStatusConditionIfNotMatch(t *testing.T) {
 		assert.Contains(t, err.Error(), "get error")
 		mockClient.AssertExpectations(t)
 	})
+
+	t.Run("Should update when observedGeneration does not match current generation", func(t *testing.T) {
+		mockClient := test.NewClient()
+
+		inferenceset := &kaitov1alpha1.InferenceSet{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:       "test-inferenceset",
+				Namespace:  "default",
+				Generation: 4, // Current generation
+			},
+			Status: kaitov1alpha1.InferenceSetStatus{
+				Conditions: []metav1.Condition{
+					{
+						Type:               string(kaitov1alpha1.ConditionTypeResourceStatus),
+						Status:             metav1.ConditionTrue,
+						Reason:             "ResourcesReady",
+						Message:            "All resources are ready",
+						ObservedGeneration: 1, // Stale generation
+					},
+				},
+			},
+		}
+
+		// Mock the Get call for UpdateInferenceSetStatus
+		mockClient.On("Get", mock.IsType(context.Background()),
+			client.ObjectKey{Name: "test-inferenceset", Namespace: "default"},
+			mock.IsType(&kaitov1alpha1.InferenceSet{}), mock.Anything).Run(func(args mock.Arguments) {
+			is := args.Get(2).(*kaitov1alpha1.InferenceSet)
+			*is = *inferenceset
+		}).Return(nil)
+
+		// Mock the Status().Update call
+		mockClient.StatusMock.On("Update", mock.IsType(context.Background()),
+			mock.IsType(&kaitov1alpha1.InferenceSet{}), mock.Anything).Run(func(args mock.Arguments) {
+			is := args.Get(1).(*kaitov1alpha1.InferenceSet)
+			// Verify the condition was updated with new observedGeneration
+			condition := meta.FindStatusCondition(is.Status.Conditions, string(kaitov1alpha1.ConditionTypeResourceStatus))
+			assert.NotNil(t, condition)
+			assert.Equal(t, int64(4), condition.ObservedGeneration)
+		}).Return(nil)
+
+		ctx := context.Background()
+		err := UpdateStatusConditionIfNotMatch(ctx, mockClient, inferenceset,
+			kaitov1alpha1.ConditionTypeResourceStatus, metav1.ConditionTrue, "ResourcesReady", "All resources are ready")
+
+		assert.NoError(t, err)
+		mockClient.AssertExpectations(t)
+		mockClient.StatusMock.AssertExpectations(t)
+	})
 }
 
 func TestUpdateInferenceSetStatus(t *testing.T) {

pkg/utils/workspace/workspace.go

@@ -93,7 +93,7 @@ var (
 func UpdateStatusConditionIfNotMatch(ctx context.Context, c client.Client, wObj *kaitov1beta1.Workspace, cType kaitov1beta1.ConditionType,
 	cStatus metav1.ConditionStatus, cReason, cMessage string) error {
 	if curCondition := meta.FindStatusCondition(wObj.Status.Conditions, string(cType)); curCondition != nil {
-		if curCondition.Status == cStatus && curCondition.Reason == cReason && curCondition.Message == cMessage {
+		if curCondition.Status == cStatus && curCondition.Reason == cReason && curCondition.Message == cMessage && curCondition.ObservedGeneration == wObj.GetGeneration() {
 			// Nothing to change
 			return nil
 		}

pkg/utils/workspace/workspace_test.go

@@ -43,10 +43,11 @@ func TestUpdateStatusConditionIfNotMatch(t *testing.T) {
 			Status: kaitov1beta1.WorkspaceStatus{
 				Conditions: []metav1.Condition{
 					{
-						Type:    string(kaitov1beta1.ConditionTypeResourceStatus),
-						Status:  metav1.ConditionTrue,
-						Reason:  "ResourcesReady",
-						Message: "All resources are ready",
+						Type:               string(kaitov1beta1.ConditionTypeResourceStatus),
+						Status:             metav1.ConditionTrue,
+						Reason:             "ResourcesReady",
+						Message:            "All resources are ready",
+						ObservedGeneration: 1,
 					},
 				},
 			},
@@ -272,6 +273,55 @@ func TestUpdateStatusConditionIfNotMatch(t *testing.T) {
 		assert.Contains(t, err.Error(), "get error")
 		mockClient.AssertExpectations(t)
 	})
+
+	t.Run("Should update when observedGeneration does not match current generation", func(t *testing.T) {
+		mockClient := test.NewClient()
+
+		workspace := &kaitov1beta1.Workspace{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:       "test-workspace",
+				Namespace:  "default",
+				Generation: 4, // Current generation
+			},
+			Status: kaitov1beta1.WorkspaceStatus{
+				Conditions: []metav1.Condition{
+					{
+						Type:               string(kaitov1beta1.ConditionTypeResourceStatus),
+						Status:             metav1.ConditionTrue,
+						Reason:             "ResourcesReady",
+						Message:            "All resources are ready",
+						ObservedGeneration: 1, // Stale generation
+					},
+				},
+			},
+		}
+
+		// Mock the Get call for UpdateWorkspaceStatus
+		mockClient.On("Get", mock.IsType(context.Background()),
+			client.ObjectKey{Name: "test-workspace", Namespace: "default"},
+			mock.IsType(&kaitov1beta1.Workspace{}), mock.Anything).Run(func(args mock.Arguments) {
+			ws := args.Get(2).(*kaitov1beta1.Workspace)
+			*ws = *workspace
+		}).Return(nil)
+
+		// Mock the Status().Update call
+		mockClient.StatusMock.On("Update", mock.IsType(context.Background()),
+			mock.IsType(&kaitov1beta1.Workspace{}), mock.Anything).Run(func(args mock.Arguments) {
+			ws := args.Get(1).(*kaitov1beta1.Workspace)
+			// Verify the condition was updated with new observedGeneration
+			condition := meta.FindStatusCondition(ws.Status.Conditions, string(kaitov1beta1.ConditionTypeResourceStatus))
+			assert.NotNil(t, condition)
+			assert.Equal(t, int64(4), condition.ObservedGeneration)
+		}).Return(nil)
+
+		ctx := context.Background()
+		err := UpdateStatusConditionIfNotMatch(ctx, mockClient, workspace,
+			kaitov1beta1.ConditionTypeResourceStatus, metav1.ConditionTrue, "ResourcesReady", "All resources are ready")
+
+		assert.NoError(t, err)
+		mockClient.AssertExpectations(t)
+		mockClient.StatusMock.AssertExpectations(t)
+	})
 }
 
 func TestUpdateWorkspaceStatus(t *testing.T) {