fix order of mock method calls

Signed-off-by: tariq-hasan <[email protected]>
kubeflow · Aug 18, 2024 · c8b7f49 · c8b7f49
1 parent 52c55c6
commit c8b7f49
Showing 1 changed file with 9 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
@@ -241,13 +222,15 @@ spec:
 	}
 
 	cases := map[string]struct {
+		mockConfigMapGetter            *gomock.Call
 		instance                       *experimentsv1beta1.Experiment
 		parameterAssignments           []commonapiv1beta1.ParameterAssignment
 		wantRunSpecWithHyperParameters *unstructured.Unstructured
 		wantError                      error
 	}{
-		// validGetConfigMap1 case
 		"Run with valid parameters": {
+			mockConfigMapGetter: c.EXPECT().GetConfigMap(gomock.Any(), gomock.Any()).Return(
+				map[string]string{templatePath: trialSpec}, nil),
 			instance: func() *experimentsv1beta1.Experiment {
 				i := newFakeInstance()
 				i.Spec.TrialTemplate.TrialSource = experimentsv1beta1.TrialSource{
@@ -262,8 +245,9 @@ spec:
 			parameterAssignments:           newFakeParameterAssignment(),
 			wantRunSpecWithHyperParameters: expectedRunSpec,
 		},
-		// invalidConfigMapName case
 		"Invalid ConfigMap name": {
+			mockConfigMapGetter: c.EXPECT().GetConfigMap(gomock.Any(), gomock.Any()).Return(
+				nil, errConfigMapNotFound),
 			instance: func() *experimentsv1beta1.Experiment {
 				i := newFakeInstance()
 				i.Spec.TrialTemplate.TrialSource = experimentsv1beta1.TrialSource{
@@ -276,8 +260,9 @@ spec:
 			parameterAssignments: newFakeParameterAssignment(),
 			wantError:            errConfigMapNotFound,
 		},
-		// validGetConfigMap3 case
 		"Invalid template path in ConfigMap name": {
+			mockConfigMapGetter: c.EXPECT().GetConfigMap(gomock.Any(), gomock.Any()).Return(
+				map[string]string{templatePath: trialSpec}, nil),
 			instance: func() *experimentsv1beta1.Experiment {
 				i := newFakeInstance()
 				i.Spec.TrialTemplate.TrialSource = experimentsv1beta1.TrialSource{
@@ -292,10 +277,11 @@ spec:
 			parameterAssignments: newFakeParameterAssignment(),
 			wantError:            errTrialTemplateNotFound,
 		},
-		// invalidTemplate case
 		// Trial template is a string in ConfigMap
 		// Because of that, user can specify not valid unstructured template
 		"Invalid trial spec in ConfigMap": {
+			mockConfigMapGetter: c.EXPECT().GetConfigMap(gomock.Any(), gomock.Any()).Return(
+				map[string]string{templatePath: invalidTrialSpec}, nil),
 			instance: func() *experimentsv1beta1.Experiment {
 				i := newFakeInstance()
 				i.Spec.TrialTemplate.TrialSource = experimentsv1beta1.TrialSource{