Test async activity completion

Author: dandavison

tests/test_serialization_context.py

@@ -9,6 +9,7 @@
 from itertools import zip_longest
 from pprint import pformat, pprint
 from typing import Any, Literal, Optional, Type
+from warnings import warn
 
 from temporalio import activity, workflow
 from temporalio.api.common.v1 import Payload
@@ -314,6 +315,111 @@ async def test_workflow_payload_conversion(
             pprint(result.items)
 
 
+async_activity_started = asyncio.Event()
+
+
+# Async activity completion test
+@activity.defn
+async def async_activity() -> TraceData:
+    async_activity_started.set()
+    activity.raise_complete_async()
+
+
+@workflow.defn
+class AsyncActivityCompletionSerializationContextTestWorkflow:
+    @workflow.run
+    async def run(self) -> TraceData:
+        return await workflow.execute_activity(
+            async_activity,
+            start_to_close_timeout=timedelta(seconds=10),
+            activity_id="async-activity-id",
+        )
+
+
+async def test_async_activity_completion_payload_conversion(
+    client: Client,
+):
+    workflow_id = str(uuid.uuid4())
+    task_queue = str(uuid.uuid4())
+
+    config = client.config()
+    config["data_converter"] = data_converter
+    client = Client(**config)
+
+    async with Worker(
+        client,
+        task_queue=task_queue,
+        workflows=[AsyncActivityCompletionSerializationContextTestWorkflow],
+        activities=[async_activity],
+        workflow_runner=UnsandboxedWorkflowRunner(),  # so that we can use isinstance
+    ):
+        wf_handle = await client.start_workflow(
+            AsyncActivityCompletionSerializationContextTestWorkflow.run,
+            id=workflow_id,
+            task_queue=task_queue,
+        )
+        activity_handle = client.get_async_activity_handle(
+            workflow_id=workflow_id,
+            run_id=wf_handle.first_execution_run_id,
+            activity_id="async-activity-id",
+        )
+        await async_activity_started.wait()
+        await activity_handle.complete(TraceData())
+        result = await wf_handle.result()
+
+        workflow_context = dataclasses.asdict(
+            WorkflowSerializationContext(
+                namespace="default",
+                workflow_id=workflow_id,
+            )
+        )
+        activity_context = dataclasses.asdict(
+            ActivitySerializationContext(
+                namespace="default",
+                workflow_id=workflow_id,
+                workflow_type="AsyncActivityCompletionSerializationContextTestWorkflow",
+                activity_type="async_activity",
+                activity_task_queue=task_queue,
+                is_local=False,
+            )
+        )
+        assert_trace(
+            result.items,
+            [
+                TraceItem(
+                    context_type="activity",
+                    in_workflow=False,
+                    method="to_payload",
+                    context=activity_context,  # Outbound activity input
+                ),
+                TraceItem(
+                    context_type="activity",
+                    in_workflow=False,
+                    method="from_payload",
+                    context=activity_context,  # Inbound activity input
+                ),
+                TraceItem(
+                    context_type="activity",
+                    in_workflow=False,
+                    method="to_payload",
+                    context=activity_context,  # Outbound activity result
+                ),
+                TraceItem(
+                    context_type="activity",
+                    in_workflow=False,
+                    method="from_payload",
+                    context=activity_context,  # Inbound activity result
+                ),
+                TraceItem(
+                    context_type="workflow",
+                    in_workflow=True,
+                    method="to_payload",
+                    context=workflow_context,  # Inbound activity result
+                ),
+            ],
+        )
+
+
 # Signal test
 
 
@@ -567,9 +673,7 @@ async def test_update_payload_conversion(
 
 def assert_trace(trace: list[TraceItem], expected: list[TraceItem]):
     if len(trace) != len(expected):
-        raise AssertionError(
-            f"expected {len(expected)} trace items but received {len(trace)}"
-        )
+        warn(f"expected {len(expected)} trace items but received {len(trace)}")
     history = []
     for item, expected_item in zip_longest(trace, expected):
         if item is None: