fix: set type metadata on search attribute payloads in test server (#2677)

Author: AdityaTeltia

temporal-test-server/src/main/java/io/temporal/internal/testservice/TestVisibilityStoreImpl.java

@@ -1,6 +1,7 @@
 package io.temporal.internal.testservice;
 
 import com.google.common.collect.ImmutableMap;
+import com.google.protobuf.ByteString;
 import io.grpc.Status;
 import io.temporal.api.common.v1.Payload;
 import io.temporal.api.common.v1.SearchAttributes;
@@ -14,6 +15,7 @@
 
 class TestVisibilityStoreImpl implements TestVisibilityStore {
 
+  private static final String METADATA_TYPE_KEY = "type";
   private static final String DEFAULT_KEY_STRING = "CustomStringField";
   private static final String DEFAULT_KEY_TEXT = "CustomTextField";
   private static final String DEFAULT_KEY_KEYWORD = "CustomKeywordField";
@@ -76,13 +78,38 @@ public SearchAttributes getSearchAttributesForExecution(ExecutionId executionId)
   public SearchAttributes upsertSearchAttributesForExecution(
       ExecutionId executionId, @Nonnull SearchAttributes searchAttributes) {
     validateSearchAttributes(searchAttributes);
+
+    SearchAttributes.Builder searchAttributesWithType = SearchAttributes.newBuilder();
+    Map<String, IndexedValueType> registeredAttributes = getRegisteredSearchAttributes();
+
+    for (Map.Entry<String, Payload> entry : searchAttributes.getIndexedFieldsMap().entrySet()) {
+      String attributeName = entry.getKey();
+      Payload payload = entry.getValue();
+      IndexedValueType expectedType = registeredAttributes.get(attributeName);
+
+      if (expectedType != null && !payload.getMetadataMap().containsKey(METADATA_TYPE_KEY)) {
+        // Add type metadata if it's missing
+        payload =
+            payload.toBuilder()
+                .putMetadata(
+                    METADATA_TYPE_KEY,
+                    ByteString.copyFromUtf8(
+                        ProtoEnumNameUtils.uniqueToSimplifiedName(expectedType)))
+                .build();
+      }
+
+      searchAttributesWithType.putIndexedFields(attributeName, payload);
+    }
+
+    SearchAttributes searchAttributesWithMetadata = searchAttributesWithType.build();
+
     return executionSearchAttributes.compute(
         executionId,
         (key, value) ->
             value == null
-                ? searchAttributes
+                ? searchAttributesWithMetadata
                 : value.toBuilder()
-                    .putAllIndexedFields(searchAttributes.getIndexedFieldsMap())
+                    .putAllIndexedFields(searchAttributesWithMetadata.getIndexedFieldsMap())
                     .build());
   }
 

temporal-test-server/src/test/java/io/temporal/internal/testservice/TestVisibilityStoreImplTest.java

@@ -0,0 +1,92 @@
+package io.temporal.internal.testservice;
+
+import static org.junit.Assert.*;
+
+import com.google.protobuf.ByteString;
+import io.temporal.api.common.v1.Payload;
+import io.temporal.api.common.v1.SearchAttributes;
+import io.temporal.api.enums.v1.IndexedValueType;
+import io.temporal.internal.common.ProtoEnumNameUtils;
+import org.junit.Test;
+
+public class TestVisibilityStoreImplTest {
+  private static final String METADATA_TYPE_KEY = "type";
+  private static final String DEFAULT_KEY_INTEGER = "CustomIntField";
+
+  @Test
+  public void testTypeMetadataIsAddedToPayloads() {
+    TestVisibilityStoreImpl visibilityStore = new TestVisibilityStoreImpl();
+    Payload payloadWithoutType =
+        Payload.newBuilder().setData(ByteString.copyFromUtf8("test data")).build();
+
+    assertFalse(
+        "Payload should not have type metadata initially",
+        payloadWithoutType.getMetadataMap().containsKey(METADATA_TYPE_KEY));
+
+    SearchAttributes.Builder searchAttributesBuilder = SearchAttributes.newBuilder();
+    searchAttributesBuilder.putIndexedFields(DEFAULT_KEY_INTEGER, payloadWithoutType);
+    SearchAttributes searchAttributes = searchAttributesBuilder.build();
+    ExecutionId executionId = new ExecutionId("test-namespace", "test-workflow-id", "test-run-id");
+    try {
+      SearchAttributes result =
+          visibilityStore.upsertSearchAttributesForExecution(executionId, searchAttributes);
+      assertNotNull("Result should not be null", result);
+
+      if (result.containsIndexedFields(DEFAULT_KEY_INTEGER)) {
+        Payload resultPayload = result.getIndexedFieldsOrThrow(DEFAULT_KEY_INTEGER);
+        assertTrue(
+            "Type metadata should be present",
+            resultPayload.getMetadataMap().containsKey(METADATA_TYPE_KEY));
+
+        ByteString typeMetadata = resultPayload.getMetadataMap().get(METADATA_TYPE_KEY);
+        String typeString = typeMetadata.toStringUtf8();
+        String expectedTypeString =
+            ProtoEnumNameUtils.uniqueToSimplifiedName(IndexedValueType.INDEXED_VALUE_TYPE_INT);
+
+        assertEquals("Type metadata should match expected type", expectedTypeString, typeString);
+      }
+    } catch (Exception e) {
+      assertTrue("Should be a validation error", e.getMessage().contains("invalid value"));
+    }
+  }
+
+  @Test
+  public void testExistingTypeMetadataIsPreserved() {
+    TestVisibilityStoreImpl visibilityStore = new TestVisibilityStoreImpl();
+    Payload payloadWithType =
+        Payload.newBuilder()
+            .setData(ByteString.copyFromUtf8("test data"))
+            .putMetadata(METADATA_TYPE_KEY, ByteString.copyFromUtf8("Int"))
+            .build();
+
+    assertTrue(
+        "Payload should have type metadata initially",
+        payloadWithType.getMetadataMap().containsKey(METADATA_TYPE_KEY));
+
+    SearchAttributes.Builder searchAttributesBuilder = SearchAttributes.newBuilder();
+    searchAttributesBuilder.putIndexedFields(DEFAULT_KEY_INTEGER, payloadWithType);
+    SearchAttributes searchAttributes = searchAttributesBuilder.build();
+    ExecutionId executionId = new ExecutionId("test-namespace", "test-workflow-id", "test-run-id");
+
+    try {
+      SearchAttributes result =
+          visibilityStore.upsertSearchAttributesForExecution(executionId, searchAttributes);
+
+      assertNotNull("Result should not be null", result);
+
+      if (result.containsIndexedFields(DEFAULT_KEY_INTEGER)) {
+        Payload resultPayload = result.getIndexedFieldsOrThrow(DEFAULT_KEY_INTEGER);
+        assertTrue(
+            "Type metadata should be present",
+            resultPayload.getMetadataMap().containsKey(METADATA_TYPE_KEY));
+
+        ByteString typeMetadata = resultPayload.getMetadataMap().get(METADATA_TYPE_KEY);
+        String typeString = typeMetadata.toStringUtf8();
+
+        assertEquals("Existing type metadata should be preserved", "Int", typeString);
+      }
+    } catch (Exception e) {
+      assertTrue("Should be a validation error", e.getMessage().contains("invalid value"));
+    }
+  }
+}