fixed order of mock method calls

kubeflow · Jun 29, 2024 · 03faf2b · 03faf2b
1 parent 638236d
commit 03faf2b
Showing 1 changed file with 18 additions and 23 deletions.
diff --git a/pkg/controller.v1beta1/experiment/manifest/generator_test.go b/pkg/controller.v1beta1/experiment/manifest/generator_test.go
@@ -196,25 +196,6 @@ spec:
             - --momentum=${trialParameters.momentum}
             - --invalidParameter={'num_layers': 2, 'input_sizes': [32, 32, 3]}`
 
-	validGetConfigMap1 := c.EXPECT().GetConfigMap(gomock.Any(), gomock.Any()).Return(
-		map[string]string{templatePath: trialSpec}, nil)
-
-	invalidConfigMapName := c.EXPECT().GetConfigMap(gomock.Any(), gomock.Any()).Return(
-		nil, errConfigMapNotFound)
-
-	validGetConfigMap3 := c.EXPECT().GetConfigMap(gomock.Any(), gomock.Any()).Return(
-		map[string]string{templatePath: trialSpec}, nil)
-
-	invalidTemplate := c.EXPECT().GetConfigMap(gomock.Any(), gomock.Any()).Return(
-		map[string]string{templatePath: invalidTrialSpec}, nil)
-
-	gomock.InOrder(
-		validGetConfigMap1,
-		invalidConfigMapName,
-		validGetConfigMap3,
-		invalidTemplate,
-	)
-
 	// We can't compare structures, because in ConfigMap trialSpec is a string and creationTimestamp was not added
 	expectedStr := `apiVersion: batch/v1
 kind: Job
@@ -245,8 +226,8 @@ spec:
 		parameterAssignments           []commonapiv1beta1.ParameterAssignment
 		wantRunSpecWithHyperParameters *unstructured.Unstructured
 		wantError                      error
+		setupMock                      func()
 	}{
-		// validGetConfigMap1 case
 		"Run with valid parameters": {
 			instance: func() *experimentsv1beta1.Experiment {
 				i := newFakeInstance()
@@ -261,8 +242,11 @@ spec:
 			}(),
 			parameterAssignments:           newFakeParameterAssignment(),
 			wantRunSpecWithHyperParameters: expectedRunSpec,
+			setupMock: func() {
+				c.EXPECT().GetConfigMap(gomock.Any(), gomock.Any()).Return(
+					map[string]string{templatePath: trialSpec}, nil)
+			},
 		},
-		// invalidConfigMapName case
 		"Invalid ConfigMap name": {
 			instance: func() *experimentsv1beta1.Experiment {
 				i := newFakeInstance()
@@ -275,8 +259,11 @@ spec:
 			}(),
 			parameterAssignments: newFakeParameterAssignment(),
 			wantError:            errConfigMapNotFound,
+			setupMock: func() {
+				c.EXPECT().GetConfigMap(gomock.Any(), gomock.Any()).Return(
+					nil, errConfigMapNotFound)
+			},
 		},
-		// validGetConfigMap3 case
 		"Invalid template path in ConfigMap name": {
 			instance: func() *experimentsv1beta1.Experiment {
 				i := newFakeInstance()
@@ -291,8 +278,11 @@ spec:
 			}(),
 			parameterAssignments: newFakeParameterAssignment(),
 			wantError:            errTrialTemplateNotFound,
+			setupMock: func() {
+				c.EXPECT().GetConfigMap(gomock.Any(), gomock.Any()).Return(
+					map[string]string{templatePath: trialSpec}, nil)
+			},
 		},
-		// invalidTemplate case
 		// Trial template is a string in ConfigMap
 		// Because of that, user can specify not valid unstructured template
 		"Invalid trial spec in ConfigMap": {
@@ -309,11 +299,16 @@ spec:
 			}(),
 			parameterAssignments: newFakeParameterAssignment(),
 			wantError:            errConvertStringToUnstructuredFailed,
+			setupMock: func() {
+				c.EXPECT().GetConfigMap(gomock.Any(), gomock.Any()).Return(
+					map[string]string{templatePath: invalidTrialSpec}, nil)
+			},
 		},
 	}
 
 	for name, tc := range cases {
 		t.Run(name, func(t *testing.T) {
+			tc.setupMock()
 			got, err := p.GetRunSpecWithHyperParameters(tc.instance, "trial-name", "trial-namespace", tc.parameterAssignments)
 			if diff := cmp.Diff(tc.wantError, err, cmpopts.EquateErrors()); len(diff) != 0 {
 				t.Errorf("Unexpected error from GetRunSpecWithHyperParameters (-want,+got):\n%s", diff)