Add support for typed search attributes to the test suite (#1378)

internal/internal_event_handlers.go

@@ -448,6 +448,23 @@ func (wc *workflowEnvironmentImpl) UpsertSearchAttributes(attributes map[string]
 	return nil
 }
 
+func (wc *workflowEnvironmentImpl) UpsertTypedSearchAttributes(attributes SearchAttributes) error {
+	rawSearchAttributes, err := serializeTypedSearchAttributes(attributes.untypedValue)
+	if err != nil {
+		return err
+	}
+
+	if _, ok := rawSearchAttributes.GetIndexedFields()[TemporalChangeVersion]; ok {
+		return errors.New("TemporalChangeVersion is a reserved key that cannot be set, please use other key")
+	}
+
+	attr := make(map[string]interface{})
+	for k, v := range rawSearchAttributes.GetIndexedFields() {
+		attr[k] = v
+	}
+	return wc.UpsertSearchAttributes(attr)
+}
+
 func (wc *workflowEnvironmentImpl) updateWorkflowInfoWithSearchAttributes(attributes *commonpb.SearchAttributes) {
 	wc.workflowInfo.SearchAttributes = mergeSearchAttributes(wc.workflowInfo.SearchAttributes, attributes)
 }

internal/internal_worker_base.go

@@ -123,6 +123,7 @@ type (
 		RemoveSession(sessionID string)
 		GetContextPropagators() []ContextPropagator
 		UpsertSearchAttributes(attributes map[string]interface{}) error
+		UpsertTypedSearchAttributes(attributes SearchAttributes) error
 		UpsertMemo(memoMap map[string]interface{}) error
 		GetRegistry() *registry
 		// QueueUpdate request of type name

internal/internal_workflow_testsuite.go

@@ -2342,6 +2342,37 @@ func (env *testWorkflowEnvironmentImpl) UpsertSearchAttributes(attributes map[st
 	return err
 }
 
+func validateAndSerializeTypedSearchAttributes(searchAttributes map[SearchAttributeKey]interface{}) (*commonpb.SearchAttributes, error) {
+	if len(searchAttributes) == 0 {
+		return nil, errSearchAttributesNotSet
+	}
+
+	rawSearchAttributes, err := serializeTypedSearchAttributes(searchAttributes)
+	if err != nil {
+		return nil, err
+	}
+
+	return rawSearchAttributes, nil
+}
+
+func (env *testWorkflowEnvironmentImpl) UpsertTypedSearchAttributes(attributes SearchAttributes) error {
+	// Don't immediately return the error from validateAndSerializeTypedSearchAttributes, as we may need to call the mock
+	rawSearchAttributes, err := validateAndSerializeTypedSearchAttributes(attributes.untypedValue)
+
+	env.workflowInfo.SearchAttributes = mergeSearchAttributes(env.workflowInfo.SearchAttributes, rawSearchAttributes)
+
+	mockMethod := mockMethodForUpsertTypedSearchAttributes
+	if _, ok := env.expectedWorkflowMockCalls[mockMethod]; !ok {
+		// mock not found
+		return err
+	}
+
+	args := []interface{}{attributes}
+	env.workflowMock.MethodCalled(mockMethod, args...)
+
+	return err
+}
+
 func (env *testWorkflowEnvironmentImpl) UpsertMemo(memoMap map[string]interface{}) error {
 	memo, err := validateAndSerializeMemo(memoMap, env.dataConverter)
 

internal/internal_workflow_testsuite_test.go

@@ -1773,6 +1773,55 @@ func (s *WorkflowTestSuiteUnitTest) Test_MockUpsertSearchAttributes() {
 	// mix no-mock and mock is not support
 }
 
+func (s *WorkflowTestSuiteUnitTest) Test_MockUpsertTypedSearchAttributes() {
+	CustomIntKey := NewSearchAttributeKeyInt64("CustomIntField")
+	workflowFn := func(ctx Context) error {
+		err := UpsertTypedSearchAttributes(ctx)
+		s.Error(err)
+
+		err = UpsertTypedSearchAttributes(ctx, CustomIntKey.ValueSet(1))
+		s.NoError(err)
+
+		// Falls back to the mock.Anything mock
+		CustomIntKey2 := NewSearchAttributeKeyInt64("CustomIntField2")
+		err = UpsertTypedSearchAttributes(ctx, CustomIntKey2.ValueSet(2))
+		s.NoError(err)
+
+		sa := GetTypedSearchAttributes(ctx)
+		s.NotNil(sa)
+		val, ok := sa.GetInt64(CustomIntKey)
+		s.True(ok)
+		s.Equal(int64(1), val)
+		val, ok = sa.GetInt64(CustomIntKey2)
+		s.True(ok)
+		s.Equal(int64(2), val)
+
+		return nil
+	}
+
+	// no mock
+	env := s.NewTestWorkflowEnvironment()
+	env.RegisterWorkflow(workflowFn)
+
+	env.ExecuteWorkflow(workflowFn)
+	s.True(env.IsWorkflowCompleted())
+	s.Nil(env.GetWorkflowError())
+	env.AssertExpectations(s.T())
+
+	// has mock
+	env = s.NewTestWorkflowEnvironment()
+	env.OnUpsertTypedSearchAttributes(NewSearchAttributes()).Return(errors.New("empty")).Once()
+	env.OnUpsertTypedSearchAttributes(NewSearchAttributes(CustomIntKey.ValueSet(1))).Return(nil).Once()
+	env.OnUpsertTypedSearchAttributes(mock.Anything).Return(nil).Once()
+
+	env.ExecuteWorkflow(workflowFn)
+	s.True(env.IsWorkflowCompleted())
+	s.Nil(env.GetWorkflowError())
+	env.AssertExpectations(s.T())
+
+	// mix no-mock and mock is not support
+}
+
 func (s *WorkflowTestSuiteUnitTest) Test_MockUpsertMemo() {
 	workflowFn := func(ctx Context) error {
 		memo := map[string]interface{}{}

internal/workflow.go

@@ -1363,28 +1363,13 @@ func UpsertTypedSearchAttributes(ctx Context, attributes ...SearchAttributeUpdat
 }
 
 func (wc *workflowEnvironmentInterceptor) UpsertTypedSearchAttributes(ctx Context, attributes ...SearchAttributeUpdate) error {
-	assertNotInReadOnlyState(ctx)
-
 	sa := SearchAttributes{
 		untypedValue: make(map[SearchAttributeKey]interface{}),
 	}
 	for _, attribute := range attributes {
 		attribute(&sa)
 	}
-	rawSearchAttributes, err := serializeTypedSearchAttributes(sa.untypedValue)
-	if err != nil {
-		return err
-	}
-
-	if _, ok := rawSearchAttributes.GetIndexedFields()[TemporalChangeVersion]; ok {
-		return errors.New("TemporalChangeVersion is a reserved key that cannot be set, please use other key")
-	}
-
-	attr := make(map[string]interface{})
-	for k, v := range rawSearchAttributes.GetIndexedFields() {
-		attr[k] = v
-	}
-	return wc.env.UpsertSearchAttributes(attr)
+	return wc.env.UpsertTypedSearchAttributes(sa)
 }
 
 // UpsertMemo is used to add or update workflow memo.

internal/workflow_testsuite.go

@@ -436,6 +436,7 @@ const mockMethodForSignalExternalWorkflow = "workflow.SignalExternalWorkflow"
 const mockMethodForRequestCancelExternalWorkflow = "workflow.RequestCancelExternalWorkflow"
 const mockMethodForGetVersion = "workflow.GetVersion"
 const mockMethodForUpsertSearchAttributes = "workflow.UpsertSearchAttributes"
+const mockMethodForUpsertTypedSearchAttributes = "workflow.UpsertTypedSearchAttributes"
 const mockMethodForUpsertMemo = "workflow.UpsertMemo"
 
 // OnSignalExternalWorkflow setup a mock for sending signal to external workflow.
@@ -506,11 +507,20 @@ func (e *TestWorkflowEnvironment) OnGetVersion(changeID string, minSupported, ma
 // OnUpsertSearchAttributes setup a mock for workflow.UpsertSearchAttributes call.
 // If mock is not setup, the UpsertSearchAttributes call will only validate input attributes.
 // If mock is setup, all UpsertSearchAttributes calls in workflow have to be mocked.
+// Deprecated: use OnUpsertTypedSearchAttributes instead.
 func (e *TestWorkflowEnvironment) OnUpsertSearchAttributes(attributes interface{}) *MockCallWrapper {
 	call := e.workflowMock.On(mockMethodForUpsertSearchAttributes, attributes)
 	return e.wrapWorkflowCall(call)
 }
 
+// OnUpsertTypedSearchAttributes setup a mock for workflow.UpsertTypedSearchAttributes call.
+// If mock is not setup, the UpsertTypedSearchAttributes call will only validate input attributes.
+// If mock is setup, all UpsertTypedSearchAttributes calls in workflow have to be mocked.
+func (e *TestWorkflowEnvironment) OnUpsertTypedSearchAttributes(attributes ...interface{}) *MockCallWrapper {
+	call := e.workflowMock.On(mockMethodForUpsertTypedSearchAttributes, attributes...)
+	return e.wrapWorkflowCall(call)
+}
+
 // OnUpsertMemo setup a mock for workflow.UpsertMemo call.
 // If mock is not setup, the UpsertMemo call will only validate input attributes.
 // If mock is setup, all UpsertMemo calls in workflow have to be mocked.
@@ -926,6 +936,7 @@ func (e *TestWorkflowEnvironment) SetMemoOnStart(memo map[string]interface{}) er
 }
 
 // SetSearchAttributesOnStart sets the search attributes when start workflow.
+// Deprecated: Use SetTypedSearchAttributes instead.
 func (e *TestWorkflowEnvironment) SetSearchAttributesOnStart(searchAttributes map[string]interface{}) error {
 	attr, err := serializeUntypedSearchAttributes(searchAttributes)
 	if err != nil {
@@ -935,6 +946,16 @@ func (e *TestWorkflowEnvironment) SetSearchAttributesOnStart(searchAttributes ma
 	return nil
 }
 
+// SetTypedSearchAttributesOnStart sets the search attributes when start workflow.
+func (e *TestWorkflowEnvironment) SetTypedSearchAttributesOnStart(searchAttributes SearchAttributes) error {
+	attr, err := serializeSearchAttributes(nil, searchAttributes)
+	if err != nil {
+		return err
+	}
+	e.impl.workflowInfo.SearchAttributes = attr
+	return nil
+}
+
 // AssertExpectations  asserts that everything specified with OnActivity
 // in fact called as expected.  Calls may have occurred in any order.
 func (e *TestWorkflowEnvironment) AssertExpectations(t mock.TestingT) bool {

internal/workflow_testsuite_test.go

@@ -76,6 +76,18 @@ func TestSetSearchAttributesOnStart(t *testing.T) {
 	require.NotNil(t, env.impl.workflowInfo.SearchAttributes)
 }
 
+func TestSetTypedSearchAttributesOnStart(t *testing.T) {
+	t.Parallel()
+	testSuite := &WorkflowTestSuite{}
+	env := testSuite.NewTestWorkflowEnvironment()
+
+	CustomIntKey := NewSearchAttributeKeyInt64("CustomIntField")
+
+	err := env.SetTypedSearchAttributesOnStart(NewSearchAttributes(CustomIntKey.ValueSet(1)))
+	require.NoError(t, err)
+	require.NotNil(t, env.impl.workflowInfo.SearchAttributes)
+}
+
 func TestNoExplicitRegistrationRequired(t *testing.T) {
 	testSuite := &WorkflowTestSuite{}
 	env := testSuite.NewTestWorkflowEnvironment()