Test update

Author: dandavison

tests/test_serialization_context.py

@@ -480,6 +480,88 @@ async def test_query_payload_conversion(
         )
 
 
+# Update test
+
+
+@workflow.defn
+class UpdateSerializationContextTestWorkflow:
+    @workflow.run
+    async def run(self) -> None:
+        await asyncio.Event().wait()
+
+    @workflow.update
+    def my_update(self, input: TraceData) -> TraceData:
+        return input
+
+
+async def test_update_payload_conversion(
+    client: Client,
+):
+    workflow_id = str(uuid.uuid4())
+    task_queue = str(uuid.uuid4())
+
+    data_converter = dataclasses.replace(
+        DataConverter.default,
+        payload_converter_class=SerializationContextTestPayloadConverter,
+    )
+
+    config = client.config()
+    config["data_converter"] = data_converter
+    custom_client = Client(**config)
+
+    async with Worker(
+        custom_client,
+        task_queue=task_queue,
+        workflows=[UpdateSerializationContextTestWorkflow],
+        activities=[],
+        workflow_runner=UnsandboxedWorkflowRunner(),  # so that we can use isinstance
+    ):
+        handle = await custom_client.start_workflow(
+            UpdateSerializationContextTestWorkflow.run,
+            id=workflow_id,
+            task_queue=task_queue,
+        )
+        result = await handle.execute_update(
+            UpdateSerializationContextTestWorkflow.my_update, TraceData()
+        )
+
+        workflow_context = dataclasses.asdict(
+            WorkflowSerializationContext(
+                namespace="default",
+                workflow_id=workflow_id,
+            )
+        )
+        assert_trace(
+            result.items,
+            [
+                TraceItem(
+                    context_type="workflow",
+                    in_workflow=False,
+                    method="to_payload",
+                    context=workflow_context,  # Outbound update input
+                ),
+                TraceItem(
+                    context_type="workflow",
+                    in_workflow=True,
+                    method="from_payload",
+                    context=workflow_context,  # Inbound update input
+                ),
+                TraceItem(
+                    context_type="workflow",
+                    in_workflow=True,
+                    method="to_payload",
+                    context=workflow_context,  # Outbound update result
+                ),
+                TraceItem(
+                    context_type="workflow",
+                    in_workflow=False,
+                    method="from_payload",
+                    context=workflow_context,  # Inbound update result
+                ),
+            ],
+        )
+
+
 # Utilities