Fix custom_converter and infrequent_polling samples (temporalio#95)

cretz · web-flow · commit 68ba233c5c8e · 2023-11-14T11:26:04.000-06:00
Fixes temporalio#93 Fixes temporalio#90
diff --git a/custom_converter/shared.py b/custom_converter/shared.py
@@ -0,0 +1,66 @@
+import dataclasses
+from typing import Any, Optional, Type
+
+import temporalio.converter
+from temporalio.api.common.v1 import Payload
+from temporalio.converter import (
+    CompositePayloadConverter,
+    DefaultPayloadConverter,
+    EncodingPayloadConverter,
+)
+
+
+class GreetingInput:
+    def __init__(self, name: str) -> None:
+        self.name = name
+
+
+class GreetingOutput:
+    def __init__(self, result: str) -> None:
+        self.result = result
+
+
+class GreetingEncodingPayloadConverter(EncodingPayloadConverter):
+    @property
+    def encoding(self) -> str:
+        return "text/my-greeting-encoding"
+
+    def to_payload(self, value: Any) -> Optional[Payload]:
+        if isinstance(value, GreetingInput):
+            return Payload(
+                metadata={"encoding": self.encoding.encode(), "is_input": b"true"},
+                data=value.name.encode(),
+            )
+        elif isinstance(value, GreetingOutput):
+            return Payload(
+                metadata={"encoding": self.encoding.encode()},
+                data=value.result.encode(),
+            )
+        else:
+            return None
+
+    def from_payload(self, payload: Payload, type_hint: Optional[Type] = None) -> Any:
+        if payload.metadata.get("is_input") == b"true":
+            # Confirm proper type hint if present
+            assert not type_hint or type_hint is GreetingInput
+            return GreetingInput(payload.data.decode())
+        else:
+            assert not type_hint or type_hint is GreetingOutput
+            return GreetingOutput(payload.data.decode())
+
+
+class GreetingPayloadConverter(CompositePayloadConverter):
+    def __init__(self) -> None:
+        # Just add ours as first before the defaults
+        super().__init__(
+            GreetingEncodingPayloadConverter(),
+            # TODO(cretz): Make this list available without instantiation - https://github.com/temporalio/sdk-python/issues/139
+            *DefaultPayloadConverter().converters.values(),
+        )
+
+
+# Use the default data converter, but change the payload converter.
+greeting_data_converter = dataclasses.replace(
+    temporalio.converter.default(),
+    payload_converter_class=GreetingPayloadConverter,
+)
diff --git a/custom_converter/starter.py b/custom_converter/starter.py
@@ -1,28 +1,22 @@
 import asyncio
-import dataclasses
 
-import temporalio.converter
 from temporalio.client import Client
 
-from custom_converter.worker import (
+from custom_converter.shared import (
     GreetingInput,
     GreetingOutput,
-    GreetingPayloadConverter,
-    GreetingWorkflow,
+    greeting_data_converter,
 )
+from custom_converter.workflow import GreetingWorkflow
 
 
 async def main():
     # Connect client
     client = await Client.connect(
         "localhost:7233",
-        # Use the default data converter, but change the payload converter.
         # Without this we get:
         #   TypeError: Object of type GreetingInput is not JSON serializable
-        data_converter=dataclasses.replace(
-            temporalio.converter.default(),
-            payload_converter_class=GreetingPayloadConverter,
-        ),
+        data_converter=greeting_data_converter,
     )
 
     # Run workflow
diff --git a/custom_converter/worker.py b/custom_converter/worker.py
@@ -1,74 +1,10 @@
 import asyncio
-import dataclasses
-from typing import Any, Optional, Type
 
-import temporalio.converter
-from temporalio import workflow
-from temporalio.api.common.v1 import Payload
 from temporalio.client import Client
-from temporalio.converter import (
-    CompositePayloadConverter,
-    DefaultPayloadConverter,
-    EncodingPayloadConverter,
-)
 from temporalio.worker import Worker
 
-
-class GreetingInput:
-    def __init__(self, name: str) -> None:
-        self.name = name
-
-
-class GreetingOutput:
-    def __init__(self, result: str) -> None:
-        self.result = result
-
-
-@workflow.defn
-class GreetingWorkflow:
-    @workflow.run
-    async def run(self, input: GreetingInput) -> GreetingOutput:
-        return GreetingOutput(f"Hello, {input.name}")
-
-
-class GreetingEncodingPayloadConverter(EncodingPayloadConverter):
-    @property
-    def encoding(self) -> str:
-        return "text/my-greeting-encoding"
-
-    def to_payload(self, value: Any) -> Optional[Payload]:
-        if isinstance(value, GreetingInput):
-            return Payload(
-                metadata={"encoding": self.encoding.encode(), "is_input": b"true"},
-                data=value.name.encode(),
-            )
-        elif isinstance(value, GreetingOutput):
-            return Payload(
-                metadata={"encoding": self.encoding.encode()},
-                data=value.result.encode(),
-            )
-        else:
-            return None
-
-    def from_payload(self, payload: Payload, type_hint: Optional[Type] = None) -> Any:
-        if payload.metadata.get("is_input") == b"true":
-            # Confirm proper type hint if present
-            assert not type_hint or type_hint is GreetingInput
-            return GreetingInput(payload.data.decode())
-        else:
-            assert not type_hint or type_hint is GreetingOutput
-            return GreetingOutput(payload.data.decode())
-
-
-class GreetingPayloadConverter(CompositePayloadConverter):
-    def __init__(self) -> None:
-        # Just add ours as first before the defaults
-        super().__init__(
-            GreetingEncodingPayloadConverter(),
-            # TODO(cretz): Make this list available without instantiation - https://github.com/temporalio/sdk-python/issues/139
-            *DefaultPayloadConverter().converters.values(),
-        )
-
+from custom_converter.shared import greeting_data_converter
+from custom_converter.workflow import GreetingWorkflow
 
 interrupt_event = asyncio.Event()
 
@@ -77,13 +13,9 @@ async def main():
     # Connect client
     client = await Client.connect(
         "localhost:7233",
-        # Use the default data converter, but change the payload converter.
         # Without this, when trying to run a workflow, we get:
         #   KeyError: 'Unknown payload encoding my-greeting-encoding
-        data_converter=dataclasses.replace(
-            temporalio.converter.default(),
-            payload_converter_class=GreetingPayloadConverter,
-        ),
+        data_converter=greeting_data_converter,
     )
 
     # Run a worker for the workflow
diff --git a/custom_converter/workflow.py b/custom_converter/workflow.py
@@ -0,0 +1,11 @@
+from temporalio import workflow
+
+with workflow.unsafe.imports_passed_through():
+    from custom_converter.shared import GreetingInput, GreetingOutput
+
+
+@workflow.defn
+class GreetingWorkflow:
+    @workflow.run
+    async def run(self, input: GreetingInput) -> GreetingOutput:
+        return GreetingOutput(f"Hello, {input.name}")
diff --git a/polling/infrequent/activities.py b/polling/infrequent/activities.py
@@ -1,4 +1,3 @@
-import asyncio
 from dataclasses import dataclass
 
 from temporalio import activity
@@ -15,9 +14,6 @@ class ComposeGreetingInput:
 @activity.defn
 async def compose_greeting(input: ComposeGreetingInput) -> str:
     test_service = TestService()
-    while True:
-        try:
-            result = test_service.get_service_result(input)
-            return result
-        except Exception:
-            activity.heartbeat("Invoking activity")
+    # If this raises an exception because it's not done yet, the activity will
+    # continually be scheduled for retry
+    return await test_service.get_service_result(input)
diff --git a/tests/custom_converter/__init__.py b/tests/custom_converter/__init__.py
diff --git a/tests/custom_converter/workflow_test.py b/tests/custom_converter/workflow_test.py
@@ -0,0 +1,28 @@
+import uuid
+
+from temporalio.client import Client
+from temporalio.worker import Worker
+
+from custom_converter.shared import (
+    GreetingInput,
+    GreetingOutput,
+    greeting_data_converter,
+)
+from custom_converter.workflow import GreetingWorkflow
+
+
+async def test_workflow_with_custom_converter(client: Client):
+    # Replace data converter in client
+    new_config = client.config()
+    new_config["data_converter"] = greeting_data_converter
+    client = Client(**new_config)
+    task_queue = f"tq-{uuid.uuid4()}"
+    async with Worker(client, task_queue=task_queue, workflows=[GreetingWorkflow]):
+        result = await client.execute_workflow(
+            GreetingWorkflow.run,
+            GreetingInput("Temporal"),
+            id=f"wf-{uuid.uuid4()}",
+            task_queue=task_queue,
+        )
+    assert isinstance(result, GreetingOutput)
+    assert result.result == "Hello, Temporal"
