Fix UpdateWithStart workflow args

temporalio · Oct 22, 2024 · 5b83ab8 · 5b83ab8
1 parent 0ce1d6e
commit 5b83ab8
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 3 deletions.
diff --git a/temporal-sdk/src/main/java/io/temporal/client/UpdateWithStartWorkflowOperation.java b/temporal-sdk/src/main/java/io/temporal/client/UpdateWithStartWorkflowOperation.java
@@ -278,6 +278,8 @@ public static <R> Builder<R> newBuilder(String updateName, Class<R> resultClass,
 
   private Object[] updateArgs;
 
+  private Object[] workflowArgs;
+
   private final CompletableFuture<WorkflowUpdateHandle<R>> handle;
 
   private final Functions.Proc request;
@@ -296,7 +298,7 @@ WorkflowUpdateHandle<R> invoke(Functions.Proc workflow) {
     try {
       request.apply();
       workflow.apply();
-      stub.updateWithStart(this, this.updateArgs);
+      stub.updateWithStart(this, this.workflowArgs);
       return this.handle.get();
     } catch (InterruptedException e) {
       Thread.currentThread().interrupt();
@@ -323,8 +325,9 @@ void prepareUpdate(
             .build();
   }
 
-  void prepareStart(WorkflowStub stub) {
+  void prepareStart(WorkflowStub stub, Object[] args) {
     setStub(stub);
+    this.workflowArgs = args;
   }
 
   /** Returns the result of the update request. */
@@ -357,6 +360,9 @@ public String toString() {
     if (request != null) {
       sb.append(", request=").append(request);
     }
+    if (workflowArgs != null) {
+      sb.append(", workflowArgs=").append(Arrays.toString(workflowArgs));
+    }
     if (updateArgs != null) {
       sb.append(", updateArgs=").append(Arrays.toString(updateArgs));
     }

diff --git a/temporal-sdk/src/main/java/io/temporal/client/WorkflowInvocationHandler.java b/temporal-sdk/src/main/java/io/temporal/client/WorkflowInvocationHandler.java
@@ -490,7 +490,7 @@ public void invoke(
           throw new IllegalArgumentException(
               "Method '" + method.getName() + "' is not a WorkflowMethod");
         }
-        this.operation.prepareStart(untyped);
+        this.operation.prepareStart(untyped, args);
         state = State.UPDATE_RECEIVED;
       } else {
         throw new IllegalArgumentException(

diff --git a/temporal-sdk/src/test/java/io/temporal/workflow/updateTest/UpdateWithStartTest.java b/temporal-sdk/src/test/java/io/temporal/workflow/updateTest/UpdateWithStartTest.java
@@ -218,12 +218,19 @@ public void startVariousFuncs() throws ExecutionException, InterruptedException
         WorkflowClient.updateWithStart(stubF6::func6, "1", 2, 3, 4, 5, 6, updateOp6);
 
     assertEquals("0", handle0.getResultAsync().get());
+    assertEquals("func", WorkflowStub.fromTyped(stubF).getResult(String.class));
     assertEquals("1", handle1.getResultAsync().get());
+    assertEquals("1", WorkflowStub.fromTyped(stubF1).getResult(String.class));
     assertEquals("2", handle2.getResultAsync().get());
+    assertEquals("12", WorkflowStub.fromTyped(stubF2).getResult(String.class));
     assertEquals("3", handle3.getResultAsync().get());
+    assertEquals("123", WorkflowStub.fromTyped(stubF3).getResult(String.class));
     assertEquals("4", handle4.getResultAsync().get());
+    assertEquals("1234", WorkflowStub.fromTyped(stubF4).getResult(String.class));
     assertEquals("5", handle5.getResultAsync().get());
+    assertEquals("12345", WorkflowStub.fromTyped(stubF5).getResult(String.class));
     assertEquals("6", handle6.getResultAsync().get());
+    assertEquals("123456", WorkflowStub.fromTyped(stubF6).getResult(String.class));
   }
 
   @Test