Refactor

Author: dandavison

tests/worker/test_nexus.py

@@ -297,43 +297,9 @@ async def test_async_response(
         task_queue=task_queue,
         workflow_runner=UnsandboxedWorkflowRunner(),
     ):
-        await create_nexus_endpoint(task_queue, client)
-        operation_workflow_id = "default-workflow-id"
-
-        # Start the caller workflow and wait until it confirms the Nexus operation has started.
-        block_forever_waiting_for_cancellation = request_cancel
-        start_op = WithStartWorkflowOperation(
-            MyCallerWorkflow.run,
-            args=[
-                MyInput(
-                    response_type=AsyncResponse(
-                        operation_workflow_id,
-                        block_forever_waiting_for_cancellation,
-                        use_shorthand_defined_operation,
-                    ),
-                    start_options=nexusrpc.handler.StartOperationOptions(
-                        headers={"my-header-key": "my-header-value"},
-                    ),
-                ),
-                request_cancel,
-                task_queue,
-            ],
-            id=str(uuid.uuid4()),
-            task_queue=task_queue,
-            id_conflict_policy=WorkflowIDConflictPolicy.FAIL,
+        caller_wf_handle, handler_wf_handle = await _start_wf_and_nexus_op(
+            client, task_queue, request_cancel, use_shorthand_defined_operation
         )
-
-        await client.execute_update_with_start_workflow(
-            MyCallerWorkflow.wait_nexus_operation_started,
-            start_workflow_operation=start_op,
-        )
-        caller_wf_handle = await start_op.workflow_handle()
-
-        # check that the operation-backing workflow now exists, and that (a) the handler
-        # workflow accepted the link to the calling Nexus event, and that (b) the caller
-        # workflow NexusOperationStarted event received in return a link to the
-        # operation-backing workflow.
-        handler_wf_handle = client.get_workflow_handle(operation_workflow_id)
         # TODO(dan): race here? How do we know it hasn't been canceled already?
         handler_wf_info = await handler_wf_handle.describe()
         assert handler_wf_info.status in [
@@ -379,6 +345,55 @@ async def test_async_response(
             assert result.start_options_received_by_handler
 
 
+async def _start_wf_and_nexus_op(
+    client: Client,
+    task_queue: str,
+    request_cancel: bool,
+    use_shorthand_defined_operation: bool,
+) -> tuple[WorkflowHandle, WorkflowHandle]:
+    """
+    Start the caller workflow and wait until the Nexus operation has started.
+    """
+    await create_nexus_endpoint(task_queue, client)
+    operation_workflow_id = "default-workflow-id"
+
+    # Start the caller workflow and wait until it confirms the Nexus operation has started.
+    block_forever_waiting_for_cancellation = request_cancel
+    start_op = WithStartWorkflowOperation(
+        MyCallerWorkflow.run,
+        args=[
+            MyInput(
+                response_type=AsyncResponse(
+                    operation_workflow_id,
+                    block_forever_waiting_for_cancellation,
+                    use_shorthand_defined_operation,
+                ),
+                start_options=nexusrpc.handler.StartOperationOptions(
+                    headers={"my-header-key": "my-header-value"},
+                ),
+            ),
+            request_cancel,
+            task_queue,
+        ],
+        id=str(uuid.uuid4()),
+        task_queue=task_queue,
+        id_conflict_policy=WorkflowIDConflictPolicy.FAIL,
+    )
+
+    await client.execute_update_with_start_workflow(
+        MyCallerWorkflow.wait_nexus_operation_started,
+        start_workflow_operation=start_op,
+    )
+    caller_wf_handle = await start_op.workflow_handle()
+
+    # check that the operation-backing workflow now exists, and that (a) the handler
+    # workflow accepted the link to the calling Nexus event, and that (b) the caller
+    # workflow NexusOperationStarted event received in return a link to the
+    # operation-backing workflow.
+    handler_wf_handle = client.get_workflow_handle(operation_workflow_id)
+    return caller_wf_handle, handler_wf_handle
+
+
 def make_nexus_endpoint_name(task_queue: str) -> str:
     # Create endpoints for different task queues without name collisions.
     return f"nexus-endpoint-{task_queue}"