Hold paused activities open until assertion observes paused=true

Replace the "treat gone as success" workaround with explicit synchronization
between pause_and_assert and the test activities it pauses. The helper
registers a threading.Event keyed by activity_id before calling pause_activity;
heartbeat_activity and sync_heartbeat_activity wait on that event after
catching their pause-induced cancellation. This keeps the activity in the
pending list until the helper has observed paused=true, eliminating the race
without weakening the assertion.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: brianstrauch

tests/helpers/__init__.py

@@ -3,6 +3,7 @@
 import logging.handlers
 import queue
 import socket
+import threading
 import time
 import uuid
 from collections.abc import Awaitable, Callable, Iterator, Sequence
@@ -265,8 +266,35 @@ async def get_pending_activity_info(
     return None
 
 
+_wait_for_pause_events: dict[str, threading.Event] = {}
+
+
+def wait_for_pause_release(activity_id: str) -> None:
+    """Block a sync activity until pause_and_assert releases it.
+
+    No-op if no event is registered: pause_and_assert may finish (and clean up
+    the event) before the activity catches its cancel, since the server has
+    already committed paused=true by the time pause_activity returns.
+    """
+    event = _wait_for_pause_events.get(activity_id)
+    if event is not None:
+        event.wait()
+
+
+async def async_wait_for_pause_release(activity_id: str) -> None:
+    """Block an async activity until pause_and_assert releases it. No-op if not registered."""
+    event = _wait_for_pause_events.get(activity_id)
+    if event is not None:
+        await asyncio.get_running_loop().run_in_executor(None, event.wait)
+
+
 async def pause_and_assert(client: Client, handle: WorkflowHandle, activity_id: str):
-    """Pause the given activity and assert it becomes paused."""
+    """Pause the given activity and assert it becomes paused.
+
+    Registers an event before calling the pause API so cooperating test
+    activities (those that catch the pause-induced cancel via
+    wait_for_pause_release) hang until we have observed paused=true.
+    """
     desc = await handle.describe()
     req = PauseActivityRequest(
         namespace=client.namespace,
@@ -276,19 +304,19 @@ async def pause_and_assert(client: Client, handle: WorkflowHandle, activity_id:
         ),
         id=activity_id,
     )
-    await client.workflow_service.pause_activity(req)
-
-    # Assert eventually paused. If the activity catches the pause-induced
-    # cancellation and returns, it leaves the pending list before we poll —
-    # treat that as success since these test activities only stop running
-    # because of the cancellation we triggered.
-    async def check_paused() -> None:
-        info = await get_pending_activity_info(handle, activity_id)
-        if info is None:
-            return
-        assert info.paused, f"Activity {activity_id} not yet paused"
-
-    await assert_eventually(check_paused)
+
+    _wait_for_pause_events[activity_id] = threading.Event()
+    try:
+        await client.workflow_service.pause_activity(req)
+
+        async def check_paused() -> None:
+            info = await assert_pending_activity_exists_eventually(handle, activity_id)
+            assert info.paused, f"Activity {activity_id} not yet paused"
+
+        await assert_eventually(check_paused)
+    finally:
+        _wait_for_pause_events[activity_id].set()
+        del _wait_for_pause_events[activity_id]
 
 
 async def unpause_and_assert(client: Client, handle: WorkflowHandle, activity_id: str):

tests/worker/test_workflow.py

@@ -130,12 +130,14 @@
     assert_pending_activity_exists_eventually,
     assert_task_fail_eventually,
     assert_workflow_exists_eventually,
+    async_wait_for_pause_release,
     ensure_search_attributes_present,
     find_free_port,
     get_pending_activity_info,
     new_worker,
     pause_and_assert,
     unpause_and_assert,
+    wait_for_pause_release,
     workflow_update_exists,
 )
 from tests.helpers.cache_eviction import (
@@ -7782,6 +7784,7 @@ async def heartbeat_activity(
     except (CancelledError, asyncio.CancelledError) as err:
         if not catch_err:
             raise err
+        await async_wait_for_pause_release(activity.info().activity_id)
         return activity.cancellation_details()
     finally:
         activity.heartbeat("finally-complete")
@@ -7801,6 +7804,7 @@ def sync_heartbeat_activity(
     except (CancelledError, asyncio.CancelledError) as err:
         if not catch_err:
             raise err
+        wait_for_pause_release(activity.info().activity_id)
         return activity.cancellation_details()
     finally:
         activity.heartbeat("finally-complete")