Fix mock error in OnUpsertSearchAttributes (#1017)

Author: vancexu

internal/error_test.go

@@ -58,7 +58,7 @@ var (
 	testErrorDetails4 = testStruct2{"a string", 321, &[]string{"eat", "code"}}
 )
 
-func (tes *testErrorStruct) Error() string{
+func (tes *testErrorStruct) Error() string {
 	return tes.message
 }
 

internal/internal_worker_test.go

@@ -136,9 +136,9 @@ func (s *internalWorkerTestSuite) TearDownTest() {
 
 func (s *internalWorkerTestSuite) createLocalActivityMarkerDataForTest(activityID, activityType string) []byte {
 	lamd := localActivityMarkerData{
-		ActivityID: activityID,
+		ActivityID:   activityID,
 		ActivityType: activityType,
-		ReplayTime: time.Now(),
+		ReplayTime:   time.Now(),
 	}
 
 	// encode marker data

internal/internal_workflow_testsuite.go

@@ -1849,19 +1849,21 @@ func getMockMethodForGetVersion(changeID string) string {
 }
 
 func (env *testWorkflowEnvironmentImpl) UpsertSearchAttributes(attributes map[string]interface{}) error {
-	attr, err := validateAndSerializeSearchAttributes(attributes)
-
-	env.workflowInfo.SearchAttributes = mergeSearchAttributes(env.workflowInfo.SearchAttributes, attr)
-
 	mockMethod := mockMethodForUpsertSearchAttributes
-	if _, ok := env.expectedMockCalls[mockMethod]; !ok {
-		// mock not found
-		return err
+	if _, ok := env.expectedMockCalls[mockMethod]; ok {
+		// mock found, check if return is error
+		args := []interface{}{attributes}
+		mockRet := env.mock.MethodCalled(mockMethod, args...)
+		if len(mockRet) != 1 {
+			panic(fmt.Sprintf("mock of UpsertSearchAttributes should return only one error"))
+		}
+		if mockRet[0] != nil {
+			return mockRet[0].(error)
+		}
 	}
 
-	args := []interface{}{attributes}
-	env.mock.MethodCalled(mockMethod, args...)
-
+	attr, err := validateAndSerializeSearchAttributes(attributes)
+	env.workflowInfo.SearchAttributes = mergeSearchAttributes(env.workflowInfo.SearchAttributes, attr)
 	return err
 }
 

internal/internal_workflow_testsuite_test.go

@@ -1389,6 +1389,24 @@ func (s *WorkflowTestSuiteUnitTest) Test_MockUpsertSearchAttributes() {
 	// mix no-mock and mock is not support
 }
 
+func (s *WorkflowTestSuiteUnitTest) Test_MockUpsertSearchAttributes_OnError() {
+	workflowFn := func(ctx Context) error {
+		attr := map[string]interface{}{}
+		attr["CustomIntField"] = 1
+		err := UpsertSearchAttributes(ctx, attr)
+		s.Error(err)
+		return err
+	}
+
+	env := s.NewTestWorkflowEnvironment()
+	env.OnUpsertSearchAttributes(map[string]interface{}{"CustomIntField": 1}).Return(errors.New("test error")).Once()
+
+	env.ExecuteWorkflow(workflowFn)
+	s.True(env.IsWorkflowCompleted())
+	s.Error(env.GetWorkflowError())
+	env.AssertExpectations(s.T())
+}
+
 func (s *WorkflowTestSuiteUnitTest) Test_ActivityWithThriftTypes() {
 	actualValues := []string{}
 	retVal := &shared.WorkflowExecution{WorkflowId: common.StringPtr("retwID2"), RunId: common.StringPtr("retrID2")}

internal/registry_test.go

@@ -193,7 +193,7 @@ type testWorkflowStruct struct{}
 type testActivityStruct struct{}
 
 func (ts *testWorkflowStruct) Method(ctx Context) error { return nil }
-func (ts *testActivityStruct) Method() error { return nil }
+func (ts *testActivityStruct) Method() error            { return nil }
 
 func testActivityFunction() error            { return nil }
 func testWorkflowFunction(ctx Context) error { return nil }