Clarify OnUpsertTypedSearchAttributes (#1391)

Clarify OnUpsertTypedSearchAttributes

internal/internal_search_attributes.go

@@ -287,19 +287,13 @@ func NewSearchAttributes(attributes ...SearchAttributeUpdate) SearchAttributes {
 	for _, attr := range attributes {
 		attr(&sa)
 	}
-	// remove nil values
-	for key, value := range sa.untypedValue {
-		if value == nil {
-			delete(sa.untypedValue, key)
-		}
-	}
 	return sa
 }
 
 // GetString gets a value for the given key and whether it was present.
 func (sa SearchAttributes) GetString(key SearchAttributeKeyString) (string, bool) {
 	value, ok := sa.untypedValue[key]
-	if !ok {
+	if !ok || value == nil {
 		return "", false
 	}
 	return value.(string), true
@@ -308,7 +302,7 @@ func (sa SearchAttributes) GetString(key SearchAttributeKeyString) (string, bool
 // GetKeyword gets a value for the given key and whether it was present.
 func (sa SearchAttributes) GetKeyword(key SearchAttributeKeyKeyword) (string, bool) {
 	value, ok := sa.untypedValue[key]
-	if !ok {
+	if !ok || value == nil {
 		return "", false
 	}
 	return value.(string), true
@@ -317,7 +311,7 @@ func (sa SearchAttributes) GetKeyword(key SearchAttributeKeyKeyword) (string, bo
 // GetBool gets a value for the given key and whether it was present.
 func (sa SearchAttributes) GetBool(key SearchAttributeKeyBool) (bool, bool) {
 	value, ok := sa.untypedValue[key]
-	if !ok {
+	if !ok || value == nil {
 		return false, false
 	}
 	return value.(bool), true
@@ -326,7 +320,7 @@ func (sa SearchAttributes) GetBool(key SearchAttributeKeyBool) (bool, bool) {
 // GetInt64 gets a value for the given key and whether it was present.
 func (sa SearchAttributes) GetInt64(key SearchAttributeKeyInt64) (int64, bool) {
 	value, ok := sa.untypedValue[key]
-	if !ok {
+	if !ok || value == nil {
 		return 0, false
 	}
 	return value.(int64), true
@@ -335,7 +329,7 @@ func (sa SearchAttributes) GetInt64(key SearchAttributeKeyInt64) (int64, bool) {
 // GetFloat64 gets a value for the given key and whether it was present.
 func (sa SearchAttributes) GetFloat64(key SearchAttributeKeyFloat64) (float64, bool) {
 	value, ok := sa.untypedValue[key]
-	if !ok {
+	if !ok || value == nil {
 		return 0.0, false
 	}
 	return value.(float64), true
@@ -344,7 +338,7 @@ func (sa SearchAttributes) GetFloat64(key SearchAttributeKeyFloat64) (float64, b
 // GetTime gets a value for the given key and whether it was present.
 func (sa SearchAttributes) GetTime(key SearchAttributeKeyTime) (time.Time, bool) {
 	value, ok := sa.untypedValue[key]
-	if !ok {
+	if !ok || value == nil {
 		return time.Time{}, false
 	}
 	return value.(time.Time), true
@@ -353,7 +347,7 @@ func (sa SearchAttributes) GetTime(key SearchAttributeKeyTime) (time.Time, bool)
 // GetKeywordList gets a value for the given key and whether it was present.
 func (sa SearchAttributes) GetKeywordList(key SearchAttributeKeyKeywordList) ([]string, bool) {
 	value, ok := sa.untypedValue[key]
-	if !ok {
+	if !ok || value == nil {
 		return nil, false
 	}
 	result := value.([]string)
@@ -363,19 +357,23 @@ func (sa SearchAttributes) GetKeywordList(key SearchAttributeKeyKeywordList) ([]
 
 // ContainsKey gets whether a key is present.
 func (sa SearchAttributes) ContainsKey(key SearchAttributeKey) bool {
-	_, ok := sa.untypedValue[key]
-	return ok
+	val, ok := sa.untypedValue[key]
+	return ok && val != nil
 }
 
 // Size gets the size of the attribute collection.
 func (sa SearchAttributes) Size() int {
-	return len(sa.untypedValue)
+	return len(sa.GetUntypedValues())
 }
 
 // GetUntypedValues gets a copy of the collection with raw types.
 func (sa SearchAttributes) GetUntypedValues() map[SearchAttributeKey]interface{} {
 	untypedValueCopy := make(map[SearchAttributeKey]interface{}, len(sa.untypedValue))
 	for key, value := range sa.untypedValue {
+		// Filter out nil values
+		if value == nil {
+			continue
+		}
 		switch v := value.(type) {
 		case []string:
 			untypedValueCopy[key] = append([]string(nil), v...)
@@ -389,6 +387,8 @@ func (sa SearchAttributes) GetUntypedValues() map[SearchAttributeKey]interface{}
 // Copy creates an update that copies existing values.
 func (sa SearchAttributes) Copy() SearchAttributeUpdate {
 	return func(s *SearchAttributes) {
+		// GetUntypedValues returns a copy of the map without nil values
+		// so the copy won't delete any existing values
 		untypedValues := sa.GetUntypedValues()
 		for key, value := range untypedValues {
 			s.untypedValue[key] = value

internal/internal_workflow_testsuite_test.go

@@ -1782,17 +1782,19 @@ func (s *WorkflowTestSuiteUnitTest) Test_MockUpsertTypedSearchAttributes() {
 		err = UpsertTypedSearchAttributes(ctx, CustomIntKey.ValueSet(1))
 		s.NoError(err)
 
+		err = UpsertTypedSearchAttributes(ctx, CustomIntKey.ValueUnset())
+		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)
+		_, ok := sa.GetInt64(CustomIntKey)
+		s.False(ok)
+		val, ok := sa.GetInt64(CustomIntKey2)
 		s.True(ok)
 		s.Equal(int64(2), val)
 
@@ -1812,6 +1814,7 @@ func (s *WorkflowTestSuiteUnitTest) Test_MockUpsertTypedSearchAttributes() {
 	env = s.NewTestWorkflowEnvironment()
 	env.OnUpsertTypedSearchAttributes(NewSearchAttributes()).Return(errors.New("empty")).Once()
 	env.OnUpsertTypedSearchAttributes(NewSearchAttributes(CustomIntKey.ValueSet(1))).Return(nil).Once()
+	env.OnUpsertTypedSearchAttributes(NewSearchAttributes(CustomIntKey.ValueUnset())).Return(nil).Once()
 	env.OnUpsertTypedSearchAttributes(mock.Anything).Return(nil).Once()
 
 	env.ExecuteWorkflow(workflowFn)

internal/workflow_testsuite.go

@@ -516,8 +516,9 @@ func (e *TestWorkflowEnvironment) OnUpsertSearchAttributes(attributes interface{
 // 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...)
+// Note: The mock is called with a temporal.SearchAttributes constructed from the inputs to workflow.UpsertTypedSearchAttributes.
+func (e *TestWorkflowEnvironment) OnUpsertTypedSearchAttributes(attributes interface{}) *MockCallWrapper {
+	call := e.workflowMock.On(mockMethodForUpsertTypedSearchAttributes, attributes)
 	return e.wrapWorkflowCall(call)
 }