Update query test

Author: dandavison

temporalio/worker/_workflow_instance.py

@@ -702,6 +702,7 @@ async def run_query() -> None:
                         namespace=self._info.namespace,
                         workflow_id=self._info.workflow_id,
                     )
+                    # TODO: why do we deserialize query input in workflow but not signal?
                     args = self._process_handler_args(
                         job.query_type,
                         job.arguments,

tests/test_serialization_context.py

@@ -3,7 +3,6 @@
 import asyncio
 import dataclasses
 import inspect
-import json
 import uuid
 from dataclasses import dataclass, field
 from datetime import timedelta
@@ -402,106 +401,28 @@ async def test_signal_payload_conversion_can_be_given_access_to_serialization_co
 # Query test
 
 
-@dataclass
-class QueryData:
-    query_context: Optional[WorkflowSerializationContext] = None
-    value: str = ""
-
-
 @workflow.defn
 class QuerySerializationContextTestWorkflow:
-    def __init__(self) -> None:
-        self.state = QueryData(value="workflow-state")
-
     @workflow.run
     async def run(self) -> None:
-        # Keep workflow running
-        await workflow.wait_condition(lambda: False)
+        await asyncio.Event().wait()
 
     @workflow.query
-    def my_query(self) -> QueryData:
-        return self.state
-
-
-class QuerySerializationContextTestEncodingPayloadConverter(
-    EncodingPayloadConverter, WithSerializationContext
-):
-    def __init__(self, context: Optional[SerializationContext] = None):
-        self.context = context
-
-    @property
-    def encoding(self) -> str:
-        return "test-query-serialization-context"
-
-    def with_context(
-        self, context: Optional[SerializationContext]
-    ) -> QuerySerializationContextTestEncodingPayloadConverter:
-        return QuerySerializationContextTestEncodingPayloadConverter(context)
-
-    trace = []  # Class variable to capture serialization events
-
-    def to_payload(self, value: Any) -> Optional[Payload]:
-        # Only handle QueryData objects
-        if type(value).__name__ != "QueryData":
-            return None
-
-        # Capture the context during serialization
-        if self.context and isinstance(self.context, WorkflowSerializationContext):
-            self.__class__.trace.append(
-                {
-                    "operation": "query_result_serialization",
-                    "context": self.context,
-                    "value": value.value,
-                }
-            )
-
-        # Serialize as JSON
-        data = {"value": value.value}
-        return Payload(
-            metadata={"encoding": self.encoding.encode()},
-            data=json.dumps(data).encode(),
-        )
-
-    def from_payload(self, payload: Payload, type_hint: Optional[Type] = None) -> Any:
-        data = json.loads(payload.data.decode())
-        return QueryData(
-            query_context=None,  # Context is not transmitted, only captured during serialization
-            value=data.get("value", ""),
-        )
-
-
-class QuerySerializationContextTestPayloadConverter(
-    CompositePayloadConverter, WithSerializationContext
-):
-    def __init__(self, context: Optional[SerializationContext] = None):
-        # Create converters with context
-        converters = [
-            QuerySerializationContextTestEncodingPayloadConverter(context),
-            *DefaultPayloadConverter.default_encoding_payload_converters,
-        ]
-        super().__init__(*converters)
-        self.context = context
-
-    def with_context(
-        self, context: Optional[SerializationContext]
-    ) -> QuerySerializationContextTestPayloadConverter:
-        return QuerySerializationContextTestPayloadConverter(context)
+    def my_query(self, input: TraceData) -> TraceData:
+        return input
 
 
 async def test_query_payload_conversion_can_be_given_access_to_serialization_context(
     client: Client,
 ):
     workflow_id = str(uuid.uuid4())
     task_queue = str(uuid.uuid4())
-    # Clear the trace before starting    QuerySerializationContextTestEncodingPayloadConverter.trace = []
 
-    # Create client with our custom data converter
     data_converter = dataclasses.replace(
         DataConverter.default,
-        payload_converter_class=QuerySerializationContextTestPayloadConverter,
+        payload_converter_class=SerializationContextTestPayloadConverter,
     )
 
-    # Create a new client with the custom data converter
     config = client.config()
     config["data_converter"] = data_converter
     custom_client = Client(**config)
@@ -511,34 +432,52 @@ async def test_query_payload_conversion_can_be_given_access_to_serialization_con
         task_queue=task_queue,
         workflows=[QuerySerializationContextTestWorkflow],
         activities=[],
+        workflow_runner=UnsandboxedWorkflowRunner(),  # so that we can use isinstance
     ):
-        # Start the workflow
         handle = await custom_client.start_workflow(
             QuerySerializationContextTestWorkflow.run,
             id=workflow_id,
             task_queue=task_queue,
         )
+        result = await handle.query(
+            QuerySerializationContextTestWorkflow.my_query, TraceData()
+        )
 
-        # Query the workflow
-        result = await handle.query(QuerySerializationContextTestWorkflow.my_query)
-
-        # Verify the result value
-        assert result.value == "workflow-state"
-
-        print(
-            f"DEBUG: trace length = {len(QuerySerializationContextTestEncodingPayloadConverter.trace)}"
+        workflow_context = dataclasses.asdict(
+            WorkflowSerializationContext(
+                namespace="default",
+                workflow_id=workflow_id,
+            )
         )
-        assert len(QuerySerializationContextTestEncodingPayloadConverter.trace) > 0
-        trace_entry = QuerySerializationContextTestEncodingPayloadConverter.trace[-1]
-        assert trace_entry["operation"] == "query_result_serialization"
-        assert trace_entry["context"] == 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 query input
+                ),
+                TraceItem(
+                    context_type="workflow",
+                    in_workflow=True,
+                    method="from_payload",
+                    context=workflow_context,  # Inbound query input
+                ),
+                TraceItem(
+                    context_type="workflow",
+                    in_workflow=True,
+                    method="to_payload",
+                    context=workflow_context,  # Outbound query result
+                ),
+                TraceItem(
+                    context_type="workflow",
+                    in_workflow=False,
+                    method="from_payload",
+                    context=workflow_context,  # Inbound query result
+                ),
+            ],
         )
-        assert trace_entry["value"] == "workflow-state"
-
-        # Cancel the workflow to clean up
-        await handle.cancel()
 
 
 # Utilities