From 03faf2bca54228c3212d5d89139d1e964dcd0671 Mon Sep 17 00:00:00 2001
From: tariq-hasan <mmtariquehsn@gmail.com>
Date: Sat, 29 Jun 2024 13:20:18 -0400
Subject: [PATCH] fixed order of mock method calls

---
 .../experiment/manifest/generator_test.go     | 41 ++++++++-----------
 1 file changed, 18 insertions(+), 23 deletions(-)

diff --git a/pkg/controller.v1beta1/experiment/manifest/generator_test.go b/pkg/controller.v1beta1/experiment/manifest/generator_test.go
index 6059d4d7d96..bf98dbe468e 100644
--- 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)