remove sleep calls from tests, add cancellation details to async cancellation errors from external activities

Author: THardy98

temporalio/client.py

@@ -56,6 +56,7 @@
 import temporalio.runtime
 import temporalio.service
 import temporalio.workflow
+from temporalio.activity import ActivityCancellationDetails
 from temporalio.service import (
     HttpConnectProxyConfig,
     KeepAliveConfig,
@@ -5126,9 +5127,12 @@ def __init__(self) -> None:
 class AsyncActivityCancelledError(temporalio.exceptions.TemporalError):
     """Error that occurs when async activity attempted heartbeat but was cancelled."""
 
-    def __init__(self) -> None:
+    details: Optional[ActivityCancellationDetails] = None
+
+    def __init__(self, details: Optional[ActivityCancellationDetails] = None) -> None:
         """Create async activity cancelled error."""
         super().__init__("Activity cancelled")
+        self.details = details
 
 
 class ScheduleAlreadyRunningError(temporalio.exceptions.TemporalError):
@@ -6265,7 +6269,12 @@ async def heartbeat_async_activity(
                 timeout=input.rpc_timeout,
             )
             if resp_by_id.cancel_requested or resp_by_id.activity_paused:
-                raise AsyncActivityCancelledError()
+                raise AsyncActivityCancelledError(
+                    details=ActivityCancellationDetails(
+                        cancel_requested=resp_by_id.cancel_requested,
+                        paused=resp_by_id.activity_paused,
+                    )
+                )
 
         else:
             resp = await self._client.workflow_service.record_activity_task_heartbeat(
@@ -6280,7 +6289,12 @@ async def heartbeat_async_activity(
                 timeout=input.rpc_timeout,
             )
             if resp.cancel_requested or resp.activity_paused:
-                raise AsyncActivityCancelledError()
+                raise AsyncActivityCancelledError(
+                    details=ActivityCancellationDetails(
+                        cancel_requested=resp.cancel_requested,
+                        paused=resp.activity_paused,
+                    )
+                )
 
     async def complete_async_activity(self, input: CompleteAsyncActivityInput) -> None:
         result = (

tests/helpers/__init__.py

@@ -4,7 +4,7 @@
 import uuid
 from contextlib import closing
 from datetime import timedelta
-from typing import Any, Awaitable, Callable, Optional, Sequence, Type, TypeVar, cast
+from typing import Any, Awaitable, Callable, Optional, Sequence, Type, TypeVar
 
 from temporalio.api.common.v1 import WorkflowExecution
 from temporalio.api.enums.v1 import IndexedValueType
@@ -224,17 +224,49 @@ async def assert_pending_activity_exists_eventually(
 ) -> PendingActivityInfo:
     """Wait until a pending activity with the given ID exists and return it."""
 
-    async def check() -> Optional[PendingActivityInfo]:
-        desc = await handle.describe()
-        for act in desc.raw_description.pending_activities:
-            if act.activity_id == activity_id:
-                return act
+    async def check() -> PendingActivityInfo:
+        act_info = await _get_pending_activity_info(handle, activity_id)
+        if act_info is not None:
+            return act_info
         raise AssertionError(
             f"Activity with ID {activity_id} not found in pending activities"
         )
 
-    activity_info = await assert_eventually(check, timeout=timeout)
-    return cast(PendingActivityInfo, activity_info)
+    return await assert_eventually(check, timeout=timeout)
+
+
+async def wait_for_next_heartbeat_cycle(
+    handle: WorkflowHandle,
+    activity_id: str,
+    initial_heartbeat_time: Any,
+    timeout: timedelta = timedelta(seconds=5),
+) -> None:
+    """Wait for the next heartbeat cycle by monitoring last_heartbeat_time changes."""
+
+    async def check_heartbeat_changed() -> None:
+        current_info = await _get_pending_activity_info(handle, activity_id)
+        if current_info is None:
+            raise AssertionError(
+                f"Activity with ID {activity_id} not found in pending activities"
+            )
+        if current_info.last_heartbeat_time == initial_heartbeat_time:
+            raise AssertionError(
+                f"Activity with ID {activity_id} has not heartbeated yet"
+            )
+
+    await assert_eventually(check_heartbeat_changed, timeout=timeout)
+
+
+async def _get_pending_activity_info(
+    handle: WorkflowHandle,
+    activity_id: str,
+) -> Optional[PendingActivityInfo]:
+    """Get pending activity info by ID, or None if not found."""
+    desc = await handle.describe()
+    for act in desc.raw_description.pending_activities:
+        if act.activity_id == activity_id:
+            return act
+    return None
 
 
 async def pause_and_assert(client: Client, handle: WorkflowHandle, activity_id: str):

tests/worker/test_workflow.py

@@ -49,12 +49,12 @@
 from temporalio.api.sdk.v1 import EnhancedStackTrace
 from temporalio.api.workflowservice.v1 import (
     GetWorkflowExecutionHistoryRequest,
-    PauseActivityRequest,
     ResetStickyTaskQueueRequest,
 )
 from temporalio.bridge.proto.workflow_activation import WorkflowActivation
 from temporalio.bridge.proto.workflow_completion import WorkflowActivationCompletion
 from temporalio.client import (
+    AsyncActivityCancelledError,
     Client,
     RPCError,
     RPCStatusCode,
@@ -127,6 +127,7 @@
     new_worker,
     pause_and_assert,
     unpause_and_assert,
+    wait_for_next_heartbeat_cycle,
     workflow_update_exists,
 )
 from tests.helpers.external_stack_trace import (
@@ -7637,14 +7638,16 @@ async def heartbeat_activity(
     while True:
         try:
             activity.heartbeat()
-            # If we are on the second attempt, we have retried due to pause/unpause.
-            if activity.info().attempt > 1:
+            # If we have heartbeat details, we are on the second attempt, we have retried due to pause/unpause.
+            if activity.info().heartbeat_details:
                 return activity.cancellation_details()
             await asyncio.sleep(0.1)
         except (CancelledError, asyncio.CancelledError) as err:
             if not catch_err:
                 raise err
             return activity.cancellation_details()
+        finally:
+            activity.heartbeat("finally-complete")
 
 
 @activity.defn
@@ -7654,14 +7657,16 @@ def sync_heartbeat_activity(
     while True:
         try:
             activity.heartbeat()
-            # If we are on the second attempt, we have retried due to pause/unpause.
-            if activity.info().attempt > 1:
+            # If we have heartbeat details, we are on the second attempt, we have retried due to pause/unpause.
+            if activity.info().heartbeat_details:
                 return activity.cancellation_details()
             time.sleep(0.1)
         except (CancelledError, asyncio.CancelledError) as err:
             if not catch_err:
                 raise err
             return activity.cancellation_details()
+        finally:
+            activity.heartbeat("finally-complete")
 
 
 @workflow.defn
@@ -7806,7 +7811,10 @@ async def test_activity_pause_unpause(client: Client, env: WorkflowEnvironment):
 
             # Wait for next heartbeat to propagate the cancellation. Unpausing before the heartbeat
             # will show activity as unpaused to core. Consequently, it will *not* issue an activity cancel.
-            time.sleep(0.3)
+            await wait_for_next_heartbeat_cycle(
+                handle, activity_info_1.activity_id, activity_info_1.last_heartbeat_time
+            )
+
             # Unpause activity
             await unpause_and_assert(client, handle, activity_info_1.activity_id)
             # Expect second activity to have started now
@@ -7818,11 +7826,74 @@ async def test_activity_pause_unpause(client: Client, env: WorkflowEnvironment):
             # Pause activity then assert it is paused
             await pause_and_assert(client, handle, activity_info_2.activity_id)
             # Wait for next heartbeat to propagate the cancellation.
-            time.sleep(0.3)
+            await wait_for_next_heartbeat_cycle(
+                handle, activity_info_2.activity_id, activity_info_2.last_heartbeat_time
+            )
             # Unpause activity
             await unpause_and_assert(client, handle, activity_info_2.activity_id)
 
             # Check workflow complete
             result = await handle.result()
             assert result[0] == None
             assert result[1] == None
+
+
+@activity.defn
+async def external_activity_heartbeat() -> None:
+    activity.raise_complete_async()
+
+
+@workflow.defn
+class ExternalActivityWorkflow:
+    @workflow.run
+    async def run(self, activity_id: str) -> None:
+        await workflow.execute_activity(
+            external_activity_heartbeat,
+            activity_id=activity_id,
+            start_to_close_timeout=timedelta(seconds=10),
+            heartbeat_timeout=timedelta(seconds=1),
+            retry_policy=RetryPolicy(maximum_attempts=2),
+        )
+
+
+async def test_external_activity_cancellation_details(
+    client: Client, env: WorkflowEnvironment
+):
+    if env.supports_time_skipping:
+        pytest.skip("Time-skipping server does not support pause API yet")
+    async with Worker(
+        client,
+        task_queue=str(uuid.uuid4()),
+        workflows=[ExternalActivityWorkflow],
+        activities=[external_activity_heartbeat],
+    ) as worker:
+        test_activity_id = f"heartbeat-activity-{uuid.uuid4()}"
+
+        wf_handle = await client.start_workflow(
+            ExternalActivityWorkflow.run,
+            test_activity_id,
+            id=f"test-external-activity-pause-{uuid.uuid4()}",
+            task_queue=worker.task_queue,
+        )
+        wf_desc = await wf_handle.describe()
+
+        # Wait for external activity
+        activity_info = await assert_pending_activity_exists_eventually(
+            wf_handle, test_activity_id
+        )
+        # Assert not paused
+        assert not activity_info.paused
+
+        external_activity_handle = client.get_async_activity_handle(
+            workflow_id=wf_desc.id, run_id=wf_desc.run_id, activity_id=test_activity_id
+        )
+
+        # Pause activity then assert it is paused
+        await pause_and_assert(client, wf_handle, activity_info.activity_id)
+
+        try:
+            await external_activity_handle.heartbeat()
+        except AsyncActivityCancelledError as err:
+            assert err.details == temporalio.activity.ActivityCancellationDetails(
+                paused=True
+            )