Test -- assert WAIT_COMPLETED future unblocked after handler wf cancellation

Author: dandavison

tests/nexus/test_workflow_caller_cancellation_types.py

@@ -1,8 +1,8 @@
 import asyncio
 import uuid
 from dataclasses import dataclass
-from datetime import datetime
-from typing import Optional
+from datetime import datetime, timezone
+from typing import Any, Optional
 
 import nexusrpc
 import pytest
@@ -52,6 +52,7 @@ class Input:
 @dataclass
 class CancellationResult:
     operation_token: str
+    caller_unblock_time: datetime
 
 
 @workflow.defn(sandboxed=False)
@@ -78,8 +79,12 @@ async def run(self, input: Input) -> CancellationResult:
         try:
             await op_handle
         except exceptions.NexusOperationError:
+            caller_unblock_time = workflow.now()
             assert op_handle.operation_token
-            return CancellationResult(operation_token=op_handle.operation_token)
+            return CancellationResult(
+                operation_token=op_handle.operation_token,
+                caller_unblock_time=caller_unblock_time,
+            )
         else:
             pytest.fail("Expected NexusOperationError")
 
@@ -131,12 +136,12 @@ async def check_behavior_for_try_cancel(
 
 
 async def check_behavior_for_wait_cancellation_completed(
-    caller_wf: WorkflowHandle,
+    caller_wf: WorkflowHandle[Any, CancellationResult],
     handler_wf: WorkflowHandle,
 ) -> None:
     """
-    Check that a cancellation request is sent and the caller workflow exits after the operation is
-    canceled.
+    Check that a cancellation request is sent and the caller workflow nexus operation future is
+    unblocked after the operation is canceled.
     """
     handler_status = (await handler_wf.describe()).status
     assert handler_status == WorkflowExecutionStatus.CANCELED
@@ -153,6 +158,14 @@ async def check_behavior_for_wait_cancellation_completed(
             (caller_wf, EventType.EVENT_TYPE_WORKFLOW_EXECUTION_COMPLETED),
         ]
     )
+    result = await caller_wf.result()
+    handler_wf_canceled_event_time = await _get_event_time(
+        handler_wf,
+        EventType.EVENT_TYPE_WORKFLOW_EXECUTION_CANCELED,
+    )
+    assert (
+        result.caller_unblock_time > handler_wf_canceled_event_time
+    ), "For WAIT_COMPLETED, the future should be unblocked after handler workflow cancellation. "
 
 
 @pytest.mark.parametrize(
@@ -225,6 +238,16 @@ async def _has_event(wf_handle: WorkflowHandle, event_type: EventType.ValueType)
     return False
 
 
+async def _get_event_time(
+    wf_handle: WorkflowHandle,
+    event_type: EventType.ValueType,
+) -> datetime:
+    event = await anext(
+        e async for e in wf_handle.fetch_history_events() if e.event_type == event_type
+    )
+    return event.event_time.ToDatetime().replace(tzinfo=timezone.utc)
+
+
 async def _assert_event_subsequence(
     expected_events: list[tuple[WorkflowHandle, EventType.ValueType]],
 ) -> None: