Fix test

Author: dandavison

tests/test_activity.py

@@ -67,6 +67,7 @@ async def test_get_result(client: Client):
 @dataclass
 class ActivityInput:
     wait_for_signal_workflow_id: str
+    wait_for_activity_start_workflow_id: str | None = None
 
 
 @activity.defn
@@ -189,7 +190,9 @@ async def test_manual_failure(client: Client):
             activity_id=activity_id,
             run_id=activity_handle.run_id,
         )
-        await async_activity_handle.fail(ApplicationError("Test failure"))
+        await async_activity_handle.fail(
+            ApplicationError("Test failure", non_retryable=True)
+        )
         with pytest.raises(ActivityFailedError) as err:
             await activity_handle.result()
         assert isinstance(err.value.cause, ApplicationError)
@@ -201,13 +204,22 @@ async def test_manual_failure(client: Client):
 
 @activity.defn
 async def activity_for_testing_heartbeat(input: ActivityInput) -> str:
-    wait_for_heartbeat_wf_handle = await activity.client().start_workflow(
-        WaitForSignalWorkflow.run,
-        id=input.wait_for_signal_workflow_id,
-        task_queue=activity.info().task_queue,
-    )
     info = activity.info()
     if info.attempt == 1:
+        # Signal that activity has started (only on first attempt)
+        if input.wait_for_activity_start_workflow_id:
+            await (
+                activity.client()
+                .get_workflow_handle(
+                    workflow_id=input.wait_for_activity_start_workflow_id,
+                )
+                .signal(WaitForSignalWorkflow.signal)
+            )
+        wait_for_heartbeat_wf_handle = await activity.client().start_workflow(
+            WaitForSignalWorkflow.run,
+            id=input.wait_for_signal_workflow_id,
+            task_queue=activity.info().task_queue,
+        )
         # Wait for test to notify that it has sent heartbeat
         await wait_for_heartbeat_wf_handle.result()
         raise Exception("Intentional error to force retry")
@@ -223,16 +235,25 @@ async def test_manual_heartbeat(client: Client):
     activity_id = str(uuid.uuid4())
     task_queue = str(uuid.uuid4())
     wait_for_signal_workflow_id = str(uuid.uuid4())
+    wait_for_activity_start_workflow_id = str(uuid.uuid4())
 
     activity_handle = await client.start_activity(
         activity_for_testing_heartbeat,
         args=(
-            ActivityInput(wait_for_signal_workflow_id=wait_for_signal_workflow_id),
+            ActivityInput(
+                wait_for_signal_workflow_id=wait_for_signal_workflow_id,
+                wait_for_activity_start_workflow_id=wait_for_activity_start_workflow_id,
+            ),
         ),  # TODO: overloads
         id=activity_id,
         task_queue=task_queue,
         start_to_close_timeout=timedelta(seconds=5),
     )
+    wait_for_activity_start_wf_handle = await client.start_workflow(
+        WaitForSignalWorkflow.run,
+        id=wait_for_activity_start_workflow_id,
+        task_queue=task_queue,
+    )
     async with Worker(
         client,
         task_queue=task_queue,
@@ -243,6 +264,7 @@ async def test_manual_heartbeat(client: Client):
             activity_id=activity_id,
             run_id=activity_handle.run_id,
         )
+        await wait_for_activity_start_wf_handle.result()
         await async_activity_handle.heartbeat("Test heartbeat details")
         await client.get_workflow_handle(
             workflow_id=wait_for_signal_workflow_id,