Avoid using global future

Author: dandavison

tests/nexus/test_workflow_caller_cancellation_types.py

@@ -29,9 +29,7 @@
 class TestContext:
     __test__ = False
     cancellation_type: workflow.NexusOperationCancellationType
-    caller_op_future_resolved: asyncio.Future[datetime] = field(
-        default_factory=asyncio.Future
-    )
+    caller_workflow_id: str
     cancel_handler_released: asyncio.Future[datetime] = field(
         default_factory=asyncio.Future
     )
@@ -97,7 +95,13 @@ async def cancel(
             # by the caller server. At that point, the caller server will write
             # NexusOperationCancelRequestCompleted. For TRY_CANCEL we want to prove that the nexus
             # op handle future can be resolved as cancelled before any of that.
-            await test_context.caller_op_future_resolved
+            caller_wf = nexus.client().get_workflow_handle_for(
+                CallerWorkflow.run,
+                workflow_id=test_context.caller_workflow_id,
+            )
+            await caller_wf.execute_update(
+                CallerWorkflow.wait_caller_op_future_resolved
+            )
         test_context.cancel_handler_released.set_result(datetime.now(timezone.utc))
         await super().cancel(ctx, token)
 
@@ -118,6 +122,7 @@ class Input:
 @dataclass
 class CancellationResult:
     operation_token: str
+    caller_op_future_resolved: datetime
 
 
 @workflow.defn(sandboxed=False)
@@ -130,6 +135,7 @@ def __init__(self, input: Input):
         )
         self.released = False
         self.operation_token: Optional[str] = None
+        self.caller_op_future_resolved: asyncio.Future[datetime] = asyncio.Future()
 
     @workflow.signal
     def release(self):
@@ -141,6 +147,10 @@ async def get_operation_token(self) -> str:
         assert self.operation_token
         return self.operation_token
 
+    @workflow.update
+    async def wait_caller_op_future_resolved(self) -> None:
+        await self.caller_op_future_resolved
+
     @workflow.run
     async def run(self, input: Input) -> CancellationResult:
         op_handle = await (
@@ -189,7 +199,7 @@ async def run(self, input: Input) -> CancellationResult:
         try:
             await op_handle
         except exceptions.NexusOperationError:
-            test_context.caller_op_future_resolved.set_result(workflow.now())
+            self.caller_op_future_resolved.set_result(workflow.now())
             assert op_handle.operation_token
             if input.cancellation_type in [
                 workflow.NexusOperationCancellationType.TRY_CANCEL,
@@ -209,6 +219,7 @@ async def run(self, input: Input) -> CancellationResult:
             await workflow.wait_condition(lambda: self.released)
             return CancellationResult(
                 operation_token=op_handle.operation_token,
+                caller_op_future_resolved=self.caller_op_future_resolved.result(),
             )
         else:
             pytest.fail("Expected NexusOperationError")
@@ -232,7 +243,10 @@ async def test_cancellation_type(
 
     cancellation_type = workflow.NexusOperationCancellationType[cancellation_type_name]
     global test_context
-    test_context = TestContext(cancellation_type=cancellation_type)
+    test_context = TestContext(
+        cancellation_type=cancellation_type,
+        caller_workflow_id="caller-wf-" + str(uuid.uuid4()),
+    )
 
     client = env.client
 
@@ -252,7 +266,7 @@ async def test_cancellation_type(
                 endpoint=make_nexus_endpoint_name(worker.task_queue),
                 cancellation_type=cancellation_type,
             ),
-            id="caller-wf-" + str(uuid.uuid4()),
+            id=test_context.caller_workflow_id,
             task_queue=worker.task_queue,
             id_conflict_policy=WorkflowIDConflictPolicy.FAIL,
         )
@@ -327,19 +341,18 @@ async def check_behavior_for_try_cancel(
     else:
         pytest.fail("Expected WorkflowFailureError")
     await caller_wf.signal(CallerWorkflow.release)
-    await caller_wf.result()
+    result = await caller_wf.result()
 
     handler_status = (await handler_wf.describe()).status
     assert handler_status == WorkflowExecutionStatus.CANCELED
-    caller_op_future_resolved = test_context.caller_op_future_resolved.result()
 
     await print_interleaved_histories(
         [caller_wf, handler_wf],
         extra_events=[
             (
                 caller_wf,
                 "Caller op future resolved",
-                caller_op_future_resolved,
+                result.caller_op_future_resolved,
             )
         ],
     )
@@ -354,7 +367,7 @@ async def check_behavior_for_try_cancel(
         caller_wf,
         EventType.EVENT_TYPE_NEXUS_OPERATION_CANCEL_REQUESTED,
     )
-    assert caller_op_future_resolved < op_cancel_requested_event
+    assert result.caller_op_future_resolved < op_cancel_requested_event
 
 
 async def check_behavior_for_wait_cancellation_requested(
@@ -375,7 +388,7 @@ async def check_behavior_for_wait_cancellation_requested(
         pytest.fail("Expected WorkflowFailureError")
 
     await caller_wf.signal(CallerWorkflow.release)
-    await caller_wf.result()
+    result = await caller_wf.result()
 
     handler_status = (await handler_wf.describe()).status
     assert handler_status == WorkflowExecutionStatus.CANCELED
@@ -386,7 +399,6 @@ async def check_behavior_for_wait_cancellation_requested(
             (caller_wf, EventType.EVENT_TYPE_NEXUS_OPERATION_CANCELED),
         ]
     )
-    caller_op_future_resolved = test_context.caller_op_future_resolved.result()
     op_cancel_request_completed = await get_event_time(
         caller_wf,
         EventType.EVENT_TYPE_NEXUS_OPERATION_CANCEL_REQUEST_COMPLETED,
@@ -395,7 +407,7 @@ async def check_behavior_for_wait_cancellation_requested(
         handler_wf,
         EventType.EVENT_TYPE_WORKFLOW_EXECUTION_CANCELED,
     )
-    assert op_cancel_request_completed < caller_op_future_resolved < op_canceled
+    assert op_cancel_request_completed < result.caller_op_future_resolved < op_canceled
 
 
 async def check_behavior_for_wait_cancellation_completed(
@@ -417,7 +429,7 @@ async def check_behavior_for_wait_cancellation_completed(
     assert handler_status == WorkflowExecutionStatus.CANCELED
 
     await caller_wf.signal(CallerWorkflow.release)
-    await caller_wf.result()
+    result = await caller_wf.result()
 
     await assert_event_subsequence(
         [
@@ -432,12 +444,11 @@ async def check_behavior_for_wait_cancellation_completed(
             (caller_wf, EventType.EVENT_TYPE_WORKFLOW_EXECUTION_COMPLETED),
         ]
     )
-    caller_op_future_resolved = test_context.caller_op_future_resolved.result()
     handler_wf_canceled_event_time = await get_event_time(
         handler_wf,
         EventType.EVENT_TYPE_WORKFLOW_EXECUTION_CANCELED,
     )
-    assert caller_op_future_resolved > handler_wf_canceled_event_time
+    assert result.caller_op_future_resolved > handler_wf_canceled_event_time
 
 
 async def has_event(wf_handle: WorkflowHandle, event_type: EventType.ValueType):

tests/nexus/test_workflow_caller_cancellation_types_when_cancel_handler_fails.py

@@ -36,9 +36,6 @@
 class TestContext:
     __test__ = False
     cancellation_type: workflow.NexusOperationCancellationType
-    caller_op_future_resolved: asyncio.Future[datetime] = field(
-        default_factory=asyncio.Future
-    )
     cancel_handler_released: asyncio.Future[datetime] = field(
         default_factory=asyncio.Future
     )
@@ -130,6 +127,7 @@ class Input:
 @dataclass
 class CancellationResult:
     operation_token: str
+    caller_op_future_resolved: datetime
     error_type: Optional[str] = None
     error_cause_type: Optional[str] = None
 
@@ -144,6 +142,7 @@ def __init__(self, input: Input):
         )
         self.released = False
         self.operation_token: Optional[str] = None
+        self.caller_op_future_resolved: asyncio.Future[datetime] = asyncio.Future()
 
     @workflow.signal
     def release(self):
@@ -185,13 +184,14 @@ async def run(self, input: Input) -> CancellationResult:
             error_type = err.__class__.__name__
             error_cause_type = err.__cause__.__class__.__name__
 
-        test_context.caller_op_future_resolved.set_result(workflow.now())
+        self.caller_op_future_resolved.set_result(workflow.now())
         assert op_handle.operation_token
         await workflow.wait_condition(lambda: self.released)
         return CancellationResult(
             operation_token=op_handle.operation_token,
             error_type=error_type,
             error_cause_type=error_cause_type,
+            caller_op_future_resolved=self.caller_op_future_resolved.result(),
         )
 
 
@@ -301,14 +301,13 @@ async def check_behavior_for_try_cancel(
     assert result.error_type == "NexusOperationError"
     assert result.error_cause_type == "CancelledError"
 
-    caller_op_future_resolved = test_context.caller_op_future_resolved.result()
     await print_interleaved_histories(
         [caller_wf, handler_wf],
         extra_events=[
             (
                 caller_wf,
                 "Caller op future resolved",
-                caller_op_future_resolved,
+                result.caller_op_future_resolved,
             )
         ],
     )
@@ -329,7 +328,7 @@ async def check_behavior_for_try_cancel(
         EventType.EVENT_TYPE_NEXUS_OPERATION_CANCEL_REQUEST_FAILED,
     )
     assert (
-        caller_op_future_resolved
+        result.caller_op_future_resolved
         < op_cancel_requested_event
         < op_cancel_request_failed_event
     )
@@ -353,7 +352,6 @@ async def check_behavior_for_wait_cancellation_requested(
             (caller_wf, EventType.EVENT_TYPE_WORKFLOW_EXECUTION_COMPLETED),
         ]
     )
-    caller_op_future_resolved = test_context.caller_op_future_resolved.result()
     op_cancel_request_failed = await get_event_time(
         caller_wf,
         EventType.EVENT_TYPE_NEXUS_OPERATION_CANCEL_REQUEST_FAILED,
@@ -362,7 +360,11 @@ async def check_behavior_for_wait_cancellation_requested(
         handler_wf,
         EventType.EVENT_TYPE_WORKFLOW_EXECUTION_COMPLETED,
     )
-    assert op_cancel_request_failed < caller_op_future_resolved < handler_wf_completed
+    assert (
+        op_cancel_request_failed
+        < result.caller_op_future_resolved
+        < handler_wf_completed
+    )
 
 
 async def check_behavior_for_wait_cancellation_completed(
@@ -385,9 +387,8 @@ async def check_behavior_for_wait_cancellation_completed(
             (caller_wf, EventType.EVENT_TYPE_NEXUS_OPERATION_COMPLETED),
         ]
     )
-    caller_op_future_resolved = test_context.caller_op_future_resolved.result()
     handler_wf_completed = await get_event_time(
         handler_wf,
         EventType.EVENT_TYPE_WORKFLOW_EXECUTION_COMPLETED,
     )
-    assert handler_wf_completed < caller_op_future_resolved
+    assert handler_wf_completed < result.caller_op_future_resolved