Fix test for activity reset

Author: tconley1428

temporalio/activity.py

@@ -162,6 +162,7 @@ class ActivityCancellationDetails:
     def _from_proto(
         proto: temporalio.bridge.proto.activity_task.ActivityCancellationDetails,
     ) -> ActivityCancellationDetails:
+        print(f"Building activity cancellation details, reset: {proto.is_reset}")
         return ActivityCancellationDetails(
             not_found=proto.is_not_found,
             cancel_requested=proto.is_cancelled,

temporalio/client.py

@@ -6363,12 +6363,16 @@ async def heartbeat_async_activity(
                 metadata=input.rpc_metadata,
                 timeout=input.rpc_timeout,
             )
-            if resp_by_id.cancel_requested or resp_by_id.activity_paused or resp_by_id.activity_reset:
+            if (
+                resp_by_id.cancel_requested
+                or resp_by_id.activity_paused
+                or resp_by_id.activity_reset
+            ):
                 raise AsyncActivityCancelledError(
                     details=ActivityCancellationDetails(
                         cancel_requested=resp_by_id.cancel_requested,
                         paused=resp_by_id.activity_paused,
-                        reset=resp_by_id.activity_reset
+                        reset=resp_by_id.activity_reset,
                     )
                 )
 
@@ -6389,6 +6393,7 @@ async def heartbeat_async_activity(
                     details=ActivityCancellationDetails(
                         cancel_requested=resp.cancel_requested,
                         paused=resp.activity_paused,
+                        reset=resp.activity_reset,
                     )
                 )
 

tests/worker/test_activity.py

@@ -17,7 +17,6 @@
 
 import temporalio.api.workflowservice.v1
 from temporalio import activity, workflow
-from temporalio.api.workflowservice.v1.request_response_pb2 import ResetActivityRequest
 from temporalio.client import (
     AsyncActivityHandle,
     Client,
@@ -1489,34 +1488,62 @@ async def h():
     )
     assert result.result == "details: Some detail"
 
-async def test_activity_reset(client: Client, worker: ExternalWorker):
 
+async def test_activity_reset(client: Client, worker: ExternalWorker):
     @activity.defn
     async def reset_activity() -> None:
-
-        await client.workflow_service.reset_activity(temporalio.api.workflowservice.v1.ResetActivityRequest(
+        req = temporalio.api.workflowservice.v1.ResetActivityRequest(
             namespace=client.namespace,
             execution=temporalio.api.common.v1.WorkflowExecution(
                 workflow_id=activity.info().workflow_id,
                 run_id=activity.info().workflow_run_id,
             ),
             id=activity.info().activity_id,
-            ))
+        )
+        activity.logger.info(f"Sending reset request: {req}")
+        await client.workflow_service.reset_activity(req)
         reset = False
         for _ in range(5):
+            activity.heartbeat()
             try:
                 if reset:
                     return None
-                await asyncio.sleep(1)
+                await asyncio.sleep(0.3)
             except Exception as e:
                 activity.logger.warning("Exception: ", e)
                 reset = True
                 raise
 
         assert False
 
-    await _execute_workflow_with_activity(
-        client, worker, reset_activity
-    )
+    await _execute_workflow_with_activity(client, worker, reset_activity)
+
 
+async def test_activity_reset_catch(client: Client, worker: ExternalWorker):
+    @activity.defn
+    async def wait_cancel() -> str:
+        req = temporalio.api.workflowservice.v1.ResetActivityRequest(
+            namespace=client.namespace,
+            execution=temporalio.api.common.v1.WorkflowExecution(
+                workflow_id=activity.info().workflow_id,
+                run_id=activity.info().workflow_run_id,
+            ),
+            id=activity.info().activity_id,
+        )
+        activity.logger.info(f"Sending reset request: {req}")
+        await client.workflow_service.reset_activity(req)
+        try:
+            while True:
+                await asyncio.sleep(0.3)
+                activity.heartbeat()
+        except asyncio.CancelledError:
+            details = activity.cancellation_details()
+            assert details is not None
+            return "Got cancelled error, reset? " + str(details.reset)
 
+    result = await _execute_workflow_with_activity(
+        client,
+        worker,
+        wait_cancel,
+    )
+    assert result.result == "Got cancelled error, reset? True"