Test that run ID is set in update handle

Author: dandavison

tests/worker/test_workflow.py

@@ -4425,6 +4425,46 @@ async def test_workflow_update_task_fails(client: Client, env: WorkflowEnvironme
         assert bad_validator_fail_ct == 2
 
 
+@workflow.defn
+class UpdateContinueAsNewWorkflow:
+    def __init__(self) -> None:
+        self.counter = 0
+
+    @workflow.run
+    async def run(self) -> None:
+        while True:
+            current_value = self.counter
+            await workflow.wait_condition(lambda: self.counter > current_value)
+            workflow.continue_as_new()
+
+    @workflow.update
+    async def update(self) -> None:
+        self.counter += 1
+
+
+async def test_workflow_update_handle_contains_latest_run_id(
+    client: Client, env: WorkflowEnvironment
+):
+    # The update triggers a CAN, so the first update handle returned should have the
+    # original run ID but the second should have the post-CAN run ID.
+    async with new_worker(client, UpdateContinueAsNewWorkflow) as worker:
+        wf_handle = await client.start_workflow(
+            UpdateContinueAsNewWorkflow.run,
+            id=f"update-continue-as-new-workflow-{uuid.uuid4()}",
+            task_queue=worker.task_queue,
+        )
+        upd_handle = await wf_handle.start_update(
+            UpdateContinueAsNewWorkflow.update,
+            wait_for_stage=WorkflowUpdateStage.ACCEPTED,
+        )
+        assert upd_handle.workflow_run_id == wf_handle.result_run_id
+        # upd_handle = await wf_handle.start_update(
+        #     UpdateContinueAsNewWorkflow.update,
+        #     wait_for_stage=WorkflowUpdateStage.ACCEPTED,
+        # )
+        # assert upd_handle.workflow_run_id != wf_handle.result_run_id
+
+
 @workflow.defn
 class ImmediatelyCompleteUpdateAndWorkflow:
     def __init__(self) -> None: