From 0603535842e1d6ab6d6d835e56a5bc2b6e6e16fc Mon Sep 17 00:00:00 2001
From: Rodrigo Zhou <rodrigozhou@users.noreply.github.com>
Date: Mon, 19 Sep 2022 18:26:12 -0700
Subject: [PATCH] Fix merge map of payload (#3412)

---
 common/payload/payload.go      |  7 +++++--
 common/payload/payload_test.go | 12 ++++++++++++
 host/elasticsearch_test.go     | 29 +++++++++++++++++++++++++++++
 3 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/common/payload/payload.go b/common/payload/payload.go
index 6e5ab77865c..42d0dff41b5 100644
--- a/common/payload/payload.go
+++ b/common/payload/payload.go
@@ -86,10 +86,13 @@ func MergeMapOfPayload(
 	m1 map[string]*commonpb.Payload,
 	m2 map[string]*commonpb.Payload,
 ) map[string]*commonpb.Payload {
-	if len(m1) == 0 {
-		return maps.Clone(m2)
+	if m2 == nil {
+		return maps.Clone(m1)
 	}
 	ret := maps.Clone(m1)
+	if ret == nil {
+		ret = make(map[string]*commonpb.Payload)
+	}
 	for k, v := range m2 {
 		if proto.Equal(v, nilPayload) || proto.Equal(v, emptySlicePayload) {
 			delete(ret, k)
diff --git a/common/payload/payload_test.go b/common/payload/payload_test.go
index c723f3d3692..44ce9ce17da 100644
--- a/common/payload/payload_test.go
+++ b/common/payload/payload_test.go
@@ -81,10 +81,22 @@ func TestMergeMapOfPayload(t *testing.T) {
 	resultMap := MergeMapOfPayload(currentMap, newMap)
 	assert.Equal(make(map[string]*commonpb.Payload), resultMap)
 
+	newMap = make(map[string]*commonpb.Payload)
+	resultMap = MergeMapOfPayload(currentMap, newMap)
+	assert.Equal(newMap, resultMap)
+
 	newMap = map[string]*commonpb.Payload{"key": EncodeString("val")}
 	resultMap = MergeMapOfPayload(currentMap, newMap)
 	assert.Equal(newMap, resultMap)
 
+	newMap = map[string]*commonpb.Payload{
+		"key":        EncodeString("val"),
+		"nil":        nilPayload,
+		"emptyArray": emptySlicePayload,
+	}
+	resultMap = MergeMapOfPayload(currentMap, newMap)
+	assert.Equal(map[string]*commonpb.Payload{"key": EncodeString("val")}, resultMap)
+
 	currentMap = map[string]*commonpb.Payload{"number": EncodeString("1")}
 	resultMap = MergeMapOfPayload(currentMap, newMap)
 	assert.Equal(
diff --git a/host/elasticsearch_test.go b/host/elasticsearch_test.go
index 76d964ce42c..8ef86d4a6e2 100644
--- a/host/elasticsearch_test.go
+++ b/host/elasticsearch_test.go
@@ -1333,6 +1333,35 @@ func (s *elasticsearchIntegrationSuite) TestUpsertWorkflowExecutionSearchAttribu
 		time.Sleep(waitTimeInMs * time.Millisecond)
 	}
 	s.True(verified)
+
+	// verify search attributes from DescribeWorkflowExecution
+	descRequest := &workflowservice.DescribeWorkflowExecutionRequest{
+		Namespace: s.namespace,
+		Execution: &commonpb.WorkflowExecution{
+			WorkflowId: id,
+		},
+	}
+	descResp, err := s.engine.DescribeWorkflowExecution(NewContext(), descRequest)
+	s.NoError(err)
+	expectedSearchAttributes, _ := searchattribute.Encode(
+		map[string]interface{}{
+			"CustomDoubleField":             22.0878,
+			searchattribute.BinaryChecksums: []string{"binary-v1", "binary-v2"},
+		},
+		nil,
+	)
+	s.Equal(
+		len(expectedSearchAttributes.GetIndexedFields()),
+		len(descResp.WorkflowExecutionInfo.GetSearchAttributes().GetIndexedFields()),
+	)
+	for attrName, expectedPayload := range expectedSearchAttributes.GetIndexedFields() {
+		respAttr, ok := descResp.WorkflowExecutionInfo.GetSearchAttributes().GetIndexedFields()[attrName]
+		s.True(ok)
+		s.Equal(expectedPayload.GetData(), respAttr.GetData())
+		attrType, typeSet := respAttr.GetMetadata()[searchattribute.MetadataType]
+		s.True(typeSet)
+		s.True(len(attrType) > 0)
+	}
 }
 
 func (s *elasticsearchIntegrationSuite) TestModifyWorkflowExecutionProperties() {