Clean up test

dandavison · dandavison · commit 8ac01926ba47 · 2025-04-17T08:17:07.000-04:00
diff --git a/temporalio/worker/_workflow_instance.py b/temporalio/worker/_workflow_instance.py
@@ -1821,6 +1821,20 @@ async def run_child() -> Any:
     async def _outbound_start_nexus_operation(
         self, input: StartNexusOperationInput
     ) -> _NexusOperationHandle[Any]:
+        # A Nexus operation handle contains two futures: self._start_fut is resolved wheas
+        # a result of the Nexus operation starting (activation job:
+        # resolve_nexus_operation_start), and self._result_fut is resolved as a result of
+        # the Nexus operation completing (activation job: resolve_nexus_operation). The
+        # handle itself corresponds to an asyncio.Task which waits on self.result_fut,
+        # handling CancelledError by emitting a RequestCancelNexusOperation command. We do
+        # not return the handle until we receive resolve_nexus_operation_start, like
+        # ChildWorkflowHandle and unlike ActivityHandle. Note that a Nexus operation may
+        # complete synchronously (in which case both jobs will be sent in the same
+        # activation, and start will be resolved without an operation token), or
+        # asynchronously (in which case start they may be sent in separate activations,
+        # and start will be resolved with an operation token). See comments in
+        # tests/worker/test_nexus.py for worked examples of the evolution of the resulting
+        # handle state machine in the sync and async Nexus response cases.
         handle: _NexusOperationHandle
 
         async def operation_handle_fn() -> Any:
@@ -2977,9 +2991,17 @@ async def result(self) -> O:
     def __await__(self) -> Generator[Any, Any, O]:
         return self._task.__await__()
 
+    def __repr__(self) -> str:
+        return (
+            f"{self._start_fut} "
+            f"{self._result_fut} "
+            f"Task[{self._task._state}] fut_waiter = {self._task._fut_waiter}) ({self._task._must_cancel})"
+        )
+
     def cancel(self) -> bool:
         # TODO(dan): what do we do when the start result has been delivered and we know
-        # this cannot be canceled (e.g. because it was a sync result, or because it failed.)
+        # this cannot be canceled (e.g. because it was a sync result, or because it
+        # failed.)
         return self._task.cancel()
 
     def _resolve_start_success(self, operation_token: Optional[str]) -> None:
diff --git a/tests/worker/test_nexus.py b/tests/worker/test_nexus.py
@@ -1,8 +1,7 @@
-import asyncio
 import uuid
 from dataclasses import dataclass
 from datetime import timedelta
-from typing import Union, cast
+from typing import Union
 
 import nexusrpc
 import nexusrpc.handler
@@ -34,6 +33,7 @@ class SyncResponse:
 @dataclass
 class AsyncResponse:
     operation_workflow_id: str
+    block_forever_waiting_for_cancellation: bool
 
 
 # The ordering in this union is critical since the data converter matches eagerly,
@@ -47,7 +47,6 @@ class AsyncResponse:
 @dataclass
 class MyInput:
     response_type: ResponseType
-    block_forever_waiting_for_cancellation: bool
 
 
 @dataclass
@@ -86,7 +85,7 @@ async def start(
         elif isinstance(input.response_type, AsyncResponse):
             return await temporalio.nexus.handler.start_workflow(
                 MyHandlerWorkflow.run,
-                input.block_forever_waiting_for_cancellation,
+                input.response_type.block_forever_waiting_for_cancellation,
                 id=input.response_type.operation_workflow_id,
                 options=options,
             )
@@ -126,7 +125,7 @@ class MyCallerWorkflow:
     def __init__(
         self,
         response_type: ResponseType,
-        should_cancel: bool,
+        request_cancel: bool,
         task_queue: str,
     ) -> None:
         self.nexus_service = workflow.NexusClient(
@@ -141,39 +140,33 @@ def __init__(
     async def run(
         self,
         response_type: ResponseType,
-        should_cancel: bool,
+        request_cancel: bool,
         task_queue: str,
     ) -> str:
         op_handle = await self.nexus_service.start_operation(
             MyService.my_operation,
-            MyInput(
-                response_type,
-                block_forever_waiting_for_cancellation=should_cancel,
-            ),
+            MyInput(response_type),
         )
+        print(f"🌈 {'after await start':<24}: {op_handle}")
         self._nexus_operation_started = True
-        task = cast(asyncio.Task, getattr(op_handle, "_task"))
         if isinstance(response_type, SyncResponse):
             assert op_handle.operation_token is None
-            # TODO(dan): I expected task to be done at this point
-            # assert task.done()
-            # assert not task.exception()
-            if should_cancel:
-                # TODO(dan): why does this assert pass (same Q as above re task.done())
-                assert op_handle.cancel()
-        elif isinstance(response_type, AsyncResponse):
+        else:
             assert op_handle.operation_token
-            assert not task.done()
-            # Allow the test to control when we proceed so that it can make initial
-            # assertions.
+            # Allow the test to make assertions before signalling us to proceed.
             await workflow.wait_condition(lambda: self._proceed)
+            print(f"🌈 {'after await proceed':<24}: {op_handle}")
 
-            if should_cancel:
-                # We cannot assert that cancel() returns True because it's possible that a
-                # resolve_nexus_operation job has already come in.
-                op_handle.cancel()
+        if request_cancel:
+            # Even for SyncResponse, the op_handle future is not done at this point; that
+            # transition doesn't happen until the handle is awaited.
+            print(f"🌈 {'before op_handle.cancel':<24}: {op_handle}")
+            cancel_ret = op_handle.cancel()
+            print(f"🌈 {'cancel returned':<24}: {cancel_ret}")
 
+        print(f"🌈 {'before await op_handle':<24}: {op_handle}")
         result = await op_handle
+        print(f"🌈 {'after await op_handle':<24}: {op_handle}")
         return result.val
 
     @workflow.update
@@ -190,32 +183,64 @@ def proceed(self) -> None:
 #
 
 
-# TODO(dan): cross-namespace tests
-# TODO(dan): nexus endpoint pytest fixture?
-@pytest.mark.parametrize("should_attempt_cancel", [False, True])
-async def test_sync_response(client: Client, should_attempt_cancel: bool):
-    task_queue = str(uuid.uuid4())
-    async with Worker(
-        client,
-        nexus_services=[MyServiceImpl()],
-        workflows=[MyCallerWorkflow, MyHandlerWorkflow],
-        task_queue=task_queue,
-        workflow_runner=UnsandboxedWorkflowRunner(),
-    ):
-        await create_nexus_endpoint(task_queue, client)
-        wf_handle = await client.start_workflow(
-            MyCallerWorkflow.run,
-            args=[SyncResponse(), should_attempt_cancel, task_queue],
-            id=str(uuid.uuid4()),
-            task_queue=task_queue,
-        )
-        # The response is synchronous, so the workflow's attempt to cancel the
-        # NexusOperationHandle do not result in cancellation.
-        result = await wf_handle.result()
-        assert result == "sync response"
+# When request_cancel is True, the NexusOperationHandle in the workflow evolves
+# through the following states:
+#                           start_fut             result_fut            handle_task w/ fut_waiter                        (task._must_cancel)
+#
+# Case 1: Sync Nexus operation response w/ cancellation of NexusOperationHandle
+# -----------------------------------------------------------------------------
+# >>>>>>>>>>>> WFT 1
+# after await start       : Future_7856[FINISHED] Future_7984[FINISHED] Task[PENDING] fut_waiter = Future_8240[FINISHED]) (False)
+# before op_handle.cancel : Future_7856[FINISHED] Future_7984[FINISHED] Task[PENDING] fut_waiter = Future_8240[FINISHED]) (False)
+# Future_8240[FINISHED].cancel() -> False  # no state transition; fut_waiter is already finished
+# cancel returned         : True
+# before await op_handle  : Future_7856[FINISHED] Future_7984[FINISHED] Task[PENDING] fut_waiter = Future_8240[FINISHED]) (True)
+# --> Despite cancel having been requested, this await on the nexus op handle does not
+#     raise CancelledError, because the task's underlying fut_waiter is already finished.
+# after await op_handle   : Future_7856[FINISHED] Future_7984[FINISHED] Task[FINISHED] fut_waiter = None) (False)
+# <<<<<<<<<<<< END WFT 1
+#
+
+# Case 2: Async Nexus operation response w/ cancellation of NexusOperationHandle
+# ------------------------------------------------------------------------------
+# >>>>>>>>>>>> WFT 1
+# after await start       : Future_7568[FINISHED] Future_7696[PENDING] Task[PENDING] fut_waiter = Future_7952[PENDING]) (False)
+# >>>>>>>>>>>> WFT 2
+# >>>>>>>>>>>> WFT 3
+# after await proceed     : Future_7568[FINISHED] Future_7696[PENDING] Task[PENDING] fut_waiter = Future_7952[PENDING]) (False)
+# before op_handle.cancel : Future_7568[FINISHED] Future_7696[PENDING] Task[PENDING] fut_waiter = Future_7952[PENDING]) (False)
+# Future_7952[PENDING].cancel() -> True  # transition to cancelled state; fut_waiter was not finished
+# cancel returned         : True
+# before await op_handle  : Future_7568[FINISHED] Future_7696[PENDING] Task[PENDING] fut_waiter = Future_7952[CANCELLED]) (False)
+# --> This await on the nexus op handle raises CancelledError, because the task's underlying fut_waiter is cancelled.
+
+# Thus in the sync case, although the caller workflow attempted to cancel the
+# NexusOperationHandle, this did not result in a CancelledError when the handle was
+# awaited, because both resolve_nexus_operation_start and resolve_nexus_operation jobs
+# were sent in the same activation and hence the task's fut_waiter was already finished.
+#
+# But in the async case, at the time that we cancel the NexusOperationHandle, only the
+# resolve_nexus_operation_start job had been sent; the result_fut was unresolved. Thus
+# when the handle was awaited, CancelledError was raised.
+
+# To create output like that above, set the following __repr__s:
+# asyncio.Future:
+# def __repr__(self):
+#     return f"{self.__class__.__name__}_{str(id(self))[-4:]}[{self._state}]"
+# _NexusOperationHandle:
+# def __repr__(self) -> str:
+#     return (
+#         f"{self._start_fut} "
+#         f"{self._result_fut} "
+#         f"Task[{self._task._state}] fut_waiter = {self._task._fut_waiter}) ({self._task._must_cancel})"
+#     )
 
 
-async def test_async_response(client: Client):
+# TODO(dan): cross-namespace tests
+# TODO(dan): nexus endpoint pytest fixture?
+# TODO: get rid of UnsandboxedWorkflowRunner (due to xray)
+@pytest.mark.parametrize("request_cancel", [False, True])
+async def test_sync_response(client: Client, request_cancel: bool):
     task_queue = str(uuid.uuid4())
     async with Worker(
         client,
@@ -224,37 +249,22 @@ async def test_async_response(client: Client):
         task_queue=task_queue,
         workflow_runner=UnsandboxedWorkflowRunner(),
     ):
-        operation_workflow_id = str(uuid.uuid4())
-        operation_workflow_handle = client.get_workflow_handle(operation_workflow_id)
         await create_nexus_endpoint(task_queue, client)
-
-        # Start the caller workflow
         wf_handle = await client.start_workflow(
             MyCallerWorkflow.run,
-            args=[AsyncResponse(operation_workflow_id), False, task_queue],
+            args=[SyncResponse(), request_cancel, task_queue],
             id=str(uuid.uuid4()),
             task_queue=task_queue,
         )
 
-        # Wait for the Nexus operation to start and check that the operation-backing workflow now exists.
-        await wf_handle.execute_update(MyCallerWorkflow.wait_nexus_operation_started)
-        wf_details = await operation_workflow_handle.describe()
-        assert wf_details.status in [
-            WorkflowExecutionStatus.RUNNING,
-            WorkflowExecutionStatus.COMPLETED,
-        ]
-
-        # Wait for the Nexus operation to complete and check that the operation-backing
-        # workflow has completed.
-        await wf_handle.signal(MyCallerWorkflow.proceed)
-
-        wf_details = await operation_workflow_handle.describe()
-        assert wf_details.status == WorkflowExecutionStatus.COMPLETED
+        # The operation result is returned even when request_cancel=True, because the
+        # response was synchronous and it could not be cancelled. See explanation above.
         result = await wf_handle.result()
-        assert result == "workflow result"
+        assert result == "sync response"
 
 
-async def test_cancellation_of_async_response(client: Client):
+@pytest.mark.parametrize("request_cancel", [False, True])
+async def test_async_response(client: Client, request_cancel: bool):
     task_queue = str(uuid.uuid4())
     async with Worker(
         client,
@@ -268,9 +278,16 @@ async def test_cancellation_of_async_response(client: Client):
         await create_nexus_endpoint(task_queue, client)
 
         # Start the caller workflow
+        block_forever_waiting_for_cancellation = request_cancel
         wf_handle = await client.start_workflow(
             MyCallerWorkflow.run,
-            args=[AsyncResponse(operation_workflow_id), True, task_queue],
+            args=[
+                AsyncResponse(
+                    operation_workflow_id, block_forever_waiting_for_cancellation
+                ),
+                request_cancel,
+                task_queue,
+            ],
             id=str(uuid.uuid4()),
             task_queue=task_queue,
         )
@@ -284,15 +301,23 @@ async def test_cancellation_of_async_response(client: Client):
         ]
 
         await wf_handle.signal(MyCallerWorkflow.proceed)
-        # The caller workflow will now cancel the op_handle, and await it.
 
-        # TODO(dan): assert what type of exception is raised here
-        with pytest.raises(BaseException) as ei:
-            await wf_handle.result()
-        e = ei.value
-        print(f"🌈 workflow failed: {e.__class__.__name__}({e})")
-        wf_details = await operation_workflow_handle.describe()
-        assert wf_details.status == WorkflowExecutionStatus.CANCELED
+        # The operation response was asynchronous and so request_cancel is honored. See
+        # explanation above.
+        if request_cancel:
+            # The caller workflow now cancels the op_handle, and awaits it, resulting in a
+            # CancellationError in the caller workflow.
+            with pytest.raises(BaseException) as ei:
+                await wf_handle.result()
+            e = ei.value
+            print(f"🌈 workflow failed: {e.__class__.__name__}({e})")
+            wf_details = await operation_workflow_handle.describe()
+            assert wf_details.status == WorkflowExecutionStatus.CANCELED
+        else:
+            wf_details = await operation_workflow_handle.describe()
+            assert wf_details.status == WorkflowExecutionStatus.COMPLETED
+            result = await wf_handle.result()
+            assert result == "workflow result"
 
 
 def make_nexus_endpoint_name(task_queue: str) -> str:
