Fix UpsertSearchAttributes mocking (#1027)

cadence-workflow · Oct 1, 2020 · 61f90c6 · 61f90c6
1 parent d0966cf
commit 61f90c6
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/internal/internal_workflow_testsuite.go b/internal/internal_workflow_testsuite.go
@@ -1855,10 +1855,10 @@ func (env *testWorkflowEnvironmentImpl) UpsertSearchAttributes(attributes map[st
 		// mock found, check if return is error
 		args := []interface{}{attributes}
 		mockRet := env.mock.MethodCalled(mockMethod, args...)
-		if len(mockRet) != 1 {
+		if len(mockRet) > 1 {
 			panic(fmt.Sprintf("mock of UpsertSearchAttributes should return only one error"))
 		}
-		if mockRet[0] != nil {
+		if len(mockRet) == 1 && mockRet[0] != nil {
 			return mockRet[0].(error)
 		}
 	}

diff --git a/internal/internal_workflow_testsuite_test.go b/internal/internal_workflow_testsuite_test.go
@@ -1386,6 +1386,16 @@ func (s *WorkflowTestSuiteUnitTest) Test_MockUpsertSearchAttributes() {
 	s.Nil(env.GetWorkflowError())
 	env.AssertExpectations(s.T())
 
+	// mock no return
+	env = s.NewTestWorkflowEnvironment()
+	env.OnUpsertSearchAttributes(map[string]interface{}{}).Once()
+	env.OnUpsertSearchAttributes(map[string]interface{}{"CustomIntField": 1}).Once()
+
+	env.ExecuteWorkflow(workflowFn)
+	s.True(env.IsWorkflowCompleted())
+	s.Nil(env.GetWorkflowError())
+	env.AssertExpectations(s.T())
+
 	// mix no-mock and mock is not support
 }