Support update search attribute inside workflow (#360) (#452)

Updated the TestWorkflowMutableStateImpl to ensure the UpsertWorkflowSearchAttributes is recorded to the history in the unit tests. Prior to this fix, tests that involve search attributes would
hit NonDeterminisicWorkflowError.

Co-authored-by: Liang Mei <[email protected]>

Author: emrahs

src/main/java/com/uber/cadence/internal/testservice/TestWorkflowMutableStateImpl.java

@@ -420,7 +420,8 @@ private void processDecision(
             ctx, d.getRequestCancelExternalWorkflowExecutionDecisionAttributes());
         break;
       case UpsertWorkflowSearchAttributes:
-        // TODO: https://github.com/uber/cadence-java-client/issues/360
+        processUpsertWorkflowSearchAttributes(
+            ctx, d.getUpsertWorkflowSearchAttributesDecisionAttributes(), decisionTaskCompletedId);
         break;
     }
   }
@@ -1050,6 +1051,22 @@ private void processContinueAsNewWorkflowExecution(
     event.getWorkflowExecutionContinuedAsNewEventAttributes().setNewExecutionRunId(runId);
   }
 
+  private void processUpsertWorkflowSearchAttributes(
+      RequestContext ctx,
+      UpsertWorkflowSearchAttributesDecisionAttributes attr,
+      long decisionTaskCompletedId)
+      throws BadRequestError, InternalServiceError {
+    UpsertWorkflowSearchAttributesEventAttributes upsertEventAttr =
+        new UpsertWorkflowSearchAttributesEventAttributes()
+            .setSearchAttributes(attr.getSearchAttributes())
+            .setDecisionTaskCompletedEventId(decisionTaskCompletedId);
+    HistoryEvent event =
+        new HistoryEvent()
+            .setEventType(EventType.UpsertWorkflowSearchAttributes)
+            .setUpsertWorkflowSearchAttributesEventAttributes(upsertEventAttr);
+    ctx.addEvent(event);
+  }
+
   @Override
   public void startWorkflow(
       boolean continuedAsNew, Optional<SignalWorkflowExecutionRequest> signalWithStartSignal)

src/test/java/com/uber/cadence/workflow/WorkflowTest.java

@@ -5323,13 +5323,13 @@ public void testExceptionInSignal() throws InterruptedException {
 
   public interface TestUpsertSearchAttributes {
     @WorkflowMethod
-    String execute(String keyword);
+    String execute(String taskList, String keyword);
   }
 
   public static class TestUpsertSearchAttributesImpl implements TestUpsertSearchAttributes {
 
     @Override
-    public String execute(String keyword) {
+    public String execute(String taskList, String keyword) {
       SearchAttributes searchAttributes = Workflow.getWorkflowInfo().getSearchAttributes();
       assertNull(searchAttributes);
 
@@ -5343,6 +5343,14 @@ public String execute(String keyword) {
           WorkflowUtils.getValueFromSearchAttributes(
               searchAttributes, "CustomKeywordField", String.class));
 
+      // Running the activity below ensures that we have one more decision task to be executed after
+      // adding the search attributes. This helps with replaying the history one more time to check
+      // against a possible NonDeterminisicWorkflowError which could be caused by missing
+      // UpsertWorkflowSearchAttributes event in history.
+      TestActivities activities =
+          Workflow.newActivityStub(TestActivities.class, newActivityOptions1(taskList));
+      activities.activity();
+
       return "done";
     }
   }
@@ -5353,9 +5361,9 @@ public void testUpsertSearchAttributes() {
     TestUpsertSearchAttributes testWorkflow =
         workflowClient.newWorkflowStub(
             TestUpsertSearchAttributes.class, newWorkflowOptionsBuilder(taskList).build());
-    String result = testWorkflow.execute("testKey");
+    String result = testWorkflow.execute(taskList, "testKey");
     assertEquals("done", result);
-    tracer.setExpected("upsertSearchAttributes");
+    tracer.setExpected("upsertSearchAttributes", "executeActivity TestActivities::activity");
   }
 
   private static class TracingWorkflowInterceptor implements WorkflowInterceptor {