update unit test

Signed-off-by: khhirani <[email protected]>
argoproj · Aug 6, 2021 · 0dafa0e · 0dafa0e
1 parent 4aca17c
commit 0dafa0e
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 18 deletions.
diff --git a/rollout/trafficrouting.go b/rollout/trafficrouting.go
@@ -116,7 +116,7 @@ func (c *rolloutContext) reconcileTrafficRouting() error {
 		// Checks for experiment step
 		// If current experiment exists, then create WeightDestinations for each experiment template
 		exStep := replicasetutil.GetCurrentExperimentStep(c.rollout)
-		if exStep != nil && c.currentEx != nil && c.currentEx.Status.Phase == v1alpha1.AnalysisPhaseSuccessful {
+		if exStep != nil && c.currentEx != nil && c.currentEx.Status.Phase == v1alpha1.AnalysisPhaseRunning {
 			getTemplateWeight := func(name string) *int32 {
 				for _, tmpl := range exStep.Templates {
 					if tmpl.Name == name {
@@ -127,13 +127,11 @@ func (c *rolloutContext) reconcileTrafficRouting() error {
 			}
 			for _, templateStatus := range c.currentEx.Status.TemplateStatuses {
 				templateWeight := getTemplateWeight(templateStatus.Name)
-				if templateStatus.Status == v1alpha1.TemplateStatusRunning  {
-					weightDestinations = append(weightDestinations, trafficrouting.WeightDestination{
-						ServiceName:     templateStatus.ServiceName,
-						PodTemplateHash: templateStatus.PodTemplateHash,
-						Weight:          *templateWeight,
-					})
-				}
+				weightDestinations = append(weightDestinations, trafficrouting.WeightDestination{
+					ServiceName:     templateStatus.ServiceName,
+					PodTemplateHash: templateStatus.PodTemplateHash,
+					Weight:          *templateWeight,
+				})
 			}
 		}
 	}

diff --git a/rollout/trafficrouting_test.go b/rollout/trafficrouting_test.go
@@ -215,24 +215,52 @@ func TestRolloutWithExperimentStep(t *testing.T) {
 
 	f.kubeobjects = append(f.kubeobjects, rs1, rs2, canarySvc, stableSvc)
 	f.replicaSetLister = append(f.replicaSetLister, rs1, rs2)
-	f.experimentLister = append(f.experimentLister, ex)
 
 	r2 = updateCanaryRolloutStatus(r2, rs1PodHash, 10, 0, 10, false)
 	f.rolloutLister = append(f.rolloutLister, r2)
 	f.objects = append(f.objects, r2, ex)
 
 	f.expectPatchRolloutAction(r2)
 
-	f.fakeTrafficRouting = newUnmockedFakeTrafficRoutingReconciler()
-	f.fakeTrafficRouting.On("UpdateHash", mock.Anything, mock.Anything).Return(nil)
-	f.fakeTrafficRouting.On("SetWeight", mock.Anything, mock.Anything).Return(func(desiredWeight int32, additionalDestinations ...trafficrouting.WeightDestination) error {
-		// make sure SetWeight was called with correct value
-		assert.Equal(t, int32(10), desiredWeight)
-		//assert.NotNil(t, additionalDestinations[0])
-		return nil
+	t.Run("Experiment Running - WeightDestination created", func(t *testing.T) {
+		ex.Status.Phase = v1alpha1.AnalysisPhaseRunning
+		f.fakeTrafficRouting = newUnmockedFakeTrafficRoutingReconciler()
+		f.fakeTrafficRouting.On("UpdateHash", mock.Anything, mock.Anything).Return(nil)
+		f.fakeTrafficRouting.On("SetWeight", mock.Anything, mock.Anything).Return(func(desiredWeight int32, weightDestinations ...trafficrouting.WeightDestination) error {
+			// make sure SetWeight was called with correct value
+			assert.Equal(t, int32(10), desiredWeight)
+			assert.Equal(t, int32(5), weightDestinations[0].Weight)
+			assert.Equal(t, ex.Status.TemplateStatuses[0].ServiceName, weightDestinations[0].ServiceName)
+			assert.Equal(t, ex.Status.TemplateStatuses[0].PodTemplateHash, weightDestinations[0].PodTemplateHash)
+			return nil
+		})
+		f.fakeTrafficRouting.On("VerifyWeight", mock.Anything).Return(func(desiredWeight int32, weightDestinations ...trafficrouting.WeightDestination) error {
+			assert.Equal(t, int32(10), desiredWeight)
+			assert.Equal(t, int32(5), weightDestinations[0].Weight)
+			assert.Equal(t, ex.Status.TemplateStatuses[0].ServiceName, weightDestinations[0].ServiceName)
+			assert.Equal(t, ex.Status.TemplateStatuses[0].PodTemplateHash, weightDestinations[0].PodTemplateHash)
+			return nil
+		})
+		f.run(getKey(r2, t))
+	})
+
+	t.Run("Experiment Pending - no WeightDestination created", func(t *testing.T) {
+		ex.Status.Phase = v1alpha1.AnalysisPhasePending
+		f.fakeTrafficRouting = newUnmockedFakeTrafficRoutingReconciler()
+		f.fakeTrafficRouting.On("UpdateHash", mock.Anything, mock.Anything).Return(nil)
+		f.fakeTrafficRouting.On("SetWeight", mock.Anything, mock.Anything).Return(func(desiredWeight int32, weightDestinations ...trafficrouting.WeightDestination) error {
+			// make sure SetWeight was called with correct value
+			assert.Equal(t, int32(10), desiredWeight)
+			assert.Len(t, weightDestinations, 0)
+			return nil
+		})
+		f.fakeTrafficRouting.On("VerifyWeight", mock.Anything).Return(func(desiredWeight int32, weightDestinations ...trafficrouting.WeightDestination) error {
+			assert.Equal(t, int32(10), desiredWeight)
+			assert.Len(t, weightDestinations, 0)
+			return nil
+		})
+		f.run(getKey(r2, t))
 	})
-	f.fakeTrafficRouting.On("VerifyWeight", mock.Anything).Return(true, nil)
-	f.run(getKey(r2, t))
 }
 
 func TestRolloutUsePreviousSetWeight(t *testing.T) {