update test code

Signed-off-by: Hui Kang <[email protected]>

rollout/sync_test.go

@@ -215,6 +215,54 @@ func TestReconcileRevisionHistoryLimit(t *testing.T) {
 	}
 }
 
+func TestPersistWorkloadRefGeneration(t *testing.T) {
+	replica := int32(1)
+	r := &v1alpha1.Rollout{
+		Spec: v1alpha1.RolloutSpec{
+			Replicas: &replica,
+		},
+	}
+	fake := fake.Clientset{}
+	roCtx := &rolloutContext{
+		rollout: r,
+		log:     logutil.WithRollout(r),
+		reconcilerBase: reconcilerBase{
+			argoprojclientset: &fake,
+		},
+		pauseContext: &pauseContext{
+			rollout: r,
+		},
+	}
+
+	tests := []struct {
+		annotatedRefGeneration string
+		currentObserved        string
+	}{
+		{"1", ""},
+		{"2", "1"},
+		{"", "1"},
+	}
+
+	for _, tc := range tests {
+		newStatus := &v1alpha1.RolloutStatus{
+			UpdatedReplicas:   int32(1),
+			AvailableReplicas: int32(1),
+		}
+
+		if tc.annotatedRefGeneration != "" {
+			annotations.SetRolloutWorkloadRefGeneration(r, tc.annotatedRefGeneration)
+			r.Spec.TemplateResolvedFromRef = true
+
+			newStatus.WorkloadObservedGeneration = tc.currentObserved
+		} else {
+			r.Spec.TemplateResolvedFromRef = false
+			annotations.RemoveRolloutWorkloadRefGeneration(r)
+		}
+		roCtx.persistRolloutStatus(newStatus)
+		assert.Equal(t, tc.annotatedRefGeneration, newStatus.WorkloadObservedGeneration)
+	}
+}
+
 // TestCanaryPromoteFull verifies skip pause, analysis, steps when promote full is set for a canary rollout
 func TestCanaryPromoteFull(t *testing.T) {
 	f := newFixture(t)

rollout/temlateref.go

@@ -239,22 +239,31 @@ func (r *informerBasedTemplateResolver) updateRolloutsReferenceAnnotation(obj in
 		return
 	}
 
+	var updateAnnotation func(ro *v1alpha1.Rollout)
+
 	generation := strconv.FormatInt(workloadMeta.GetGeneration(), 10)
+	updateAnnotation = func(ro *v1alpha1.Rollout) {
+		updated := annotations.SetRolloutWorkloadRefGeneration(ro, generation)
+		if updated {
+			// update the annotation causes the rollout to be requeued and the template will be resolved to the referred
+			// workload during next reconciliation
+			ro.Spec.Template.Spec.Containers = []corev1.Container{}
+			_, err := r.argoprojclientset.ArgoprojV1alpha1().Rollouts(ro.Namespace).Update(context.TODO(), ro, v1.UpdateOptions{})
+			if err != nil {
+				log.Errorf("Cannot update the workload-ref/annotation for %s/%s", ro.GetName(), ro.GetNamespace())
+			}
+		}
+	}
 	for _, ro := range rollouts {
-		un, ok := ro.(*unstructured.Unstructured)
+		rollout, ok := ro.(*v1alpha1.Rollout)
 		if ok {
-			rollout := unstructuredutil.ObjectToRollout(un)
-			updated := annotations.SetRolloutWorkloadRefGeneration(rollout, generation)
-			if updated {
-				rollout.Spec.Template.Spec.Containers = []corev1.Container{}
-				_, err := r.argoprojclientset.ArgoprojV1alpha1().Rollouts(rollout.Namespace).Update(context.TODO(), rollout, v1.UpdateOptions{})
-				if err != nil {
-					log.Errorf("Cannot update the workload-ref/annotation for %s/%s", rollout.GetName(), rollout.GetNamespace())
-				}
-			}
-
+			updateAnnotation(rollout)
 		} else {
-			fmt.Println("DEBUG: ro is not Unstructured:", ro)
+			un, ok := ro.(*unstructured.Unstructured)
+			if ok {
+				rollout := unstructuredutil.ObjectToRollout(un)
+				updateAnnotation(rollout)
+			}
 		}
 	}
 }

rollout/temlateref_test.go

@@ -2,6 +2,7 @@ package rollout
 
 import (
 	"context"
+	"fmt"
 	"testing"
 	"time"
 
@@ -22,6 +23,7 @@ import (
 	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/util/wait"
 	disco "k8s.io/client-go/discovery"
 	discofake "k8s.io/client-go/discovery/fake"
 	"k8s.io/client-go/dynamic"
@@ -377,9 +379,18 @@ func TestRequeueReferencedRollouts(t *testing.T) {
 	require.NoError(t, err)
 
 	// verify rollout's annotation is updated
-	ro, _ := rolloutsClient.ArgoprojV1alpha1().Rollouts("default").Get(context.TODO(), "my-rollout", v1.GetOptions{})
-	assert.NotNil(t, ro)
-	assert.Equal(t, "2", ro.GetAnnotations()[annotations.WorkloadGenerationAnnotation])
+	timeout := 3 * time.Second
+	err = wait.Poll(time.Second, timeout, func() (done bool, err error) {
+		ro, err := rolloutsClient.ArgoprojV1alpha1().Rollouts("default").Get(context.TODO(), "my-rollout", v1.GetOptions{})
+		if err != nil {
+			return false, fmt.Errorf("error get %v", err)
+		}
+		if "2" != ro.GetAnnotations()[annotations.WorkloadGenerationAnnotation] {
+			return false, fmt.Errorf("workload generation not equal to 2: %v", ro.GetAnnotations()[annotations.WorkloadGenerationAnnotation])
+		}
+		return true, nil
+	})
+	assert.NoError(t, err)
 
 	err = deploymentsClient.Delete(context.TODO(), deployment.Name, v1.DeleteOptions{})
 	require.NoError(t, err)

utils/annotations/annotations_test.go

@@ -152,6 +152,11 @@ func TestAnnotationUtils(t *testing.T) {
 	})
 	t.Run("RemoveRolloutWorkloadRefGeneration", func(t *testing.T) {
 		copyRollout := tRollout.DeepCopy()
+
+		copyRollout.Annotations = nil
+		RemoveRolloutWorkloadRefGeneration(copyRollout)
+		assert.Nil(t, copyRollout.Annotations)
+
 		copyRollout.Annotations = map[string]string{WorkloadGenerationAnnotation: "2"}
 		RemoveRolloutWorkloadRefGeneration(copyRollout)
 		_, ok := copyRollout.Annotations[WorkloadGenerationAnnotation]

utils/rollout/rolloututil_test.go

@@ -223,6 +223,7 @@ func TestRolloutStatusProgressing(t *testing.T) {
 	{
 		//Rollout observed workload generation is not updated
 		ro := newCanaryRollout()
+		ro.Spec.TemplateResolvedFromRef = true
 		annotations.SetRolloutWorkloadRefGeneration(ro, "2")
 		ro.Status = v1alpha1.RolloutStatus{
 			WorkloadObservedGeneration: "1",
@@ -231,6 +232,25 @@ func TestRolloutStatusProgressing(t *testing.T) {
 		assert.Equal(t, v1alpha1.RolloutPhaseProgressing, status)
 		assert.Equal(t, "waiting for rollout spec update to be observed for the reference workload", message)
 	}
+	{
+		ro := newCanaryRollout()
+
+		observed := isWorkloadGenerationObserved(ro)
+		assert.True(t, observed)
+
+		annotations.SetRolloutWorkloadRefGeneration(ro, "2")
+		ro.Status.WorkloadObservedGeneration = "222222222222222222"
+		observed = isWorkloadGenerationObserved(ro)
+		assert.True(t, observed)
+
+		ro.Status.WorkloadObservedGeneration = "1"
+		observed = isWorkloadGenerationObserved(ro)
+		assert.False(t, observed)
+
+		ro.Status.WorkloadObservedGeneration = "2"
+		observed = isWorkloadGenerationObserved(ro)
+		assert.True(t, observed)
+	}
 }
 
 func TestRolloutStatusHealthy(t *testing.T) {