Refactor test

Author: dandavison

tests/worker/test_nexus.py

@@ -2,7 +2,8 @@
 import uuid
 from dataclasses import dataclass
 from datetime import timedelta
-from typing import Optional, Union
+from enum import StrEnum
+from typing import Any, Callable, Optional, Union
 
 import nexusrpc
 import nexusrpc.handler
@@ -34,16 +35,26 @@
 # -----------------------------------------------------------------------------
 # Service interface
 #
+class CallerReference(StrEnum):
+    IMPLEMENTATION = "implementation"
+    INTERFACE = "interface"
+
+
+class OpDefinitionType(StrEnum):
+    SHORTHAND = "shorthand"
+    LONGHAND = "longhand"
+
+
 @dataclass
 class SyncResponse:
-    use_shorthand_defined_operation: bool
+    op_definition_type: OpDefinitionType
 
 
 @dataclass
 class AsyncResponse:
     operation_workflow_id: str
     block_forever_waiting_for_cancellation: bool
-    use_shorthand_defined_operation: bool
+    op_definition_type: OpDefinitionType
 
 
 # The order of the two types in this union is critical since the data converter matches
@@ -56,6 +67,7 @@ class AsyncResponse:
 class MyInput:
     response_type: ResponseType
     start_options: nexusrpc.handler.StartOperationOptions
+    caller_reference: CallerReference
 
 
 @dataclass
@@ -65,7 +77,7 @@ class MyOutput:
 
 
 @nexusrpc.interface.service
-class MyService:
+class MyServiceInterface:
     my_sync_or_async_operation: nexusrpc.interface.Operation[MyInput, MyOutput]
     my_sync_operation: nexusrpc.interface.Operation[MyInput, MyOutput]
     my_async_operation: nexusrpc.interface.Operation[MyInput, MyOutput]
@@ -127,7 +139,7 @@ async def fetch_result(self, *args, **kwargs):
         raise NotImplementedError
 
 
-@nexusrpc.handler.service(interface=MyService)
+@nexusrpc.handler.service(interface=MyServiceInterface)
 class MyServiceImpl:
     @nexusrpc.handler.operation
     def my_sync_or_async_operation(
@@ -176,7 +188,10 @@ def __init__(
         task_queue: str,
     ) -> None:
         self.nexus_service = workflow.NexusClient(
-            service=MyService,
+            service={
+                CallerReference.IMPLEMENTATION: MyServiceImpl,
+                CallerReference.INTERFACE: MyServiceInterface,
+            }[input.caller_reference],
             endpoint=make_nexus_endpoint_name(task_queue),
             schedule_to_close_timeout=timedelta(seconds=10),
         )
@@ -190,15 +205,7 @@ async def run(
         request_cancel: bool,
         task_queue: str,
     ) -> MyOutput:
-        if input.response_type.use_shorthand_defined_operation:
-            operation = (
-                MyService.my_sync_operation
-                if isinstance(input.response_type, SyncResponse)
-                else MyService.my_async_operation
-            )
-        else:
-            operation = MyService.my_sync_or_async_operation
-
+        operation = self._get_operation(input)
         op_handle = await self.nexus_service.start_operation(
             operation,
             input,
@@ -231,6 +238,61 @@ async def run(
     async def wait_nexus_operation_started(self) -> None:
         await workflow.wait_condition(lambda: self._nexus_operation_started)
 
+    def _get_operation(
+        self, input: MyInput
+    ) -> Union[
+        nexusrpc.interface.Operation[MyInput, MyOutput],
+        Callable[[Any], nexusrpc.handler.AbstractOperation[MyInput, MyOutput]],
+    ]:
+        return {
+            (
+                SyncResponse,
+                OpDefinitionType.SHORTHAND,
+                CallerReference.IMPLEMENTATION,
+            ): MyServiceImpl.my_sync_operation,
+            (
+                SyncResponse,
+                OpDefinitionType.SHORTHAND,
+                CallerReference.INTERFACE,
+            ): MyServiceInterface.my_sync_operation,
+            (
+                SyncResponse,
+                OpDefinitionType.LONGHAND,
+                CallerReference.IMPLEMENTATION,
+            ): MyServiceImpl.my_sync_or_async_operation,
+            (
+                SyncResponse,
+                OpDefinitionType.LONGHAND,
+                CallerReference.INTERFACE,
+            ): MyServiceInterface.my_sync_or_async_operation,
+            (
+                AsyncResponse,
+                OpDefinitionType.SHORTHAND,
+                CallerReference.IMPLEMENTATION,
+            ): MyServiceImpl.my_async_operation,
+            (
+                AsyncResponse,
+                OpDefinitionType.SHORTHAND,
+                CallerReference.INTERFACE,
+            ): MyServiceInterface.my_async_operation,
+            (
+                AsyncResponse,
+                OpDefinitionType.LONGHAND,
+                CallerReference.IMPLEMENTATION,
+            ): MyServiceImpl.my_sync_or_async_operation,
+            (
+                AsyncResponse,
+                OpDefinitionType.LONGHAND,
+                CallerReference.INTERFACE,
+            ): MyServiceInterface.my_sync_or_async_operation,
+        }[
+            {True: SyncResponse, False: AsyncResponse}[
+                isinstance(input.response_type, SyncResponse)
+            ],
+            input.response_type.op_definition_type,
+            input.caller_reference,
+        ]
+
 
 # -----------------------------------------------------------------------------
 # Tests
@@ -239,13 +301,20 @@ async def wait_nexus_operation_started(self) -> None:
 
 # TODO(dan): cross-namespace tests
 # TODO(dan): nexus endpoint pytest fixture?
-# TODO: get rid of UnsandboxedWorkflowRunner (due to xray)
+# TODO(dan): get rid of UnsandboxedWorkflowRunner (due to xray)
+# TODO(dan): test headers
 @pytest.mark.parametrize("request_cancel", [False, True])
-@pytest.mark.parametrize("use_shorthand_defined_operation", [False, True])
+@pytest.mark.parametrize(
+    "op_definition_type", [OpDefinitionType.SHORTHAND, OpDefinitionType.LONGHAND]
+)
+@pytest.mark.parametrize(
+    "caller_reference", [CallerReference.IMPLEMENTATION, CallerReference.INTERFACE]
+)
 async def test_sync_response(
     client: Client,
     request_cancel: bool,
-    use_shorthand_defined_operation: bool,
+    op_definition_type: OpDefinitionType,
+    caller_reference: CallerReference,
 ):
     task_queue = str(uuid.uuid4())
     async with Worker(
@@ -260,10 +329,11 @@ async def test_sync_response(
             MyCallerWorkflow.run,
             args=[
                 MyInput(
-                    response_type=SyncResponse(use_shorthand_defined_operation),
+                    response_type=SyncResponse(op_definition_type),
                     start_options=nexusrpc.handler.StartOperationOptions(
                         headers={"my-header-key": "my-header-value"},
                     ),
+                    caller_reference=caller_reference,
                 ),
                 request_cancel,
                 task_queue,
@@ -280,15 +350,19 @@ async def test_sync_response(
 
 
 @pytest.mark.parametrize("request_cancel", [False, True])
-@pytest.mark.parametrize("use_shorthand_defined_operation", [False, True])
+@pytest.mark.parametrize(
+    "op_definition_type", [OpDefinitionType.SHORTHAND, OpDefinitionType.LONGHAND]
+)
+@pytest.mark.parametrize(
+    "caller_reference", [CallerReference.IMPLEMENTATION, CallerReference.INTERFACE]
+)
 async def test_async_response(
     client: Client,
     request_cancel: bool,
-    use_shorthand_defined_operation: bool,
+    op_definition_type: OpDefinitionType,
+    caller_reference: CallerReference,
 ):
-    print(
-        f"🌈 {'test_async_response':<24}: {request_cancel=} {use_shorthand_defined_operation=}"
-    )
+    print(f"🌈 {'test_async_response':<24}: {request_cancel=} {op_definition_type=}")
     task_queue = str(uuid.uuid4())
     async with Worker(
         client,
@@ -298,7 +372,7 @@ async def test_async_response(
         workflow_runner=UnsandboxedWorkflowRunner(),
     ):
         caller_wf_handle, handler_wf_handle = await _start_wf_and_nexus_op(
-            client, task_queue, request_cancel, use_shorthand_defined_operation
+            client, task_queue, request_cancel, op_definition_type, caller_reference
         )
         # TODO(dan): race here? How do we know it hasn't been canceled already?
         handler_wf_info = await handler_wf_handle.describe()
@@ -325,10 +399,10 @@ async def test_async_response(
             # ID of first command after update accepted
             assert e.__cause__.scheduled_event_id == 6
             assert e.__cause__.endpoint == make_nexus_endpoint_name(task_queue)
-            assert e.__cause__.service == "MyService"
+            assert e.__cause__.service == "MyServiceInterface"
             assert (
                 e.__cause__.operation == "my_async_operation"
-                if use_shorthand_defined_operation
+                if op_definition_type == OpDefinitionType.SHORTHAND
                 else "my_sync_or_async_operation"
             )
             assert temporalio.nexus.handler.StartWorkflowOperationResult._decode_token(
@@ -349,7 +423,8 @@ async def _start_wf_and_nexus_op(
     client: Client,
     task_queue: str,
     request_cancel: bool,
-    use_shorthand_defined_operation: bool,
+    op_definition_type: OpDefinitionType,
+    caller_reference: CallerReference,
 ) -> tuple[WorkflowHandle, WorkflowHandle]:
     """
     Start the caller workflow and wait until the Nexus operation has started.
@@ -366,11 +441,12 @@ async def _start_wf_and_nexus_op(
                 response_type=AsyncResponse(
                     operation_workflow_id,
                     block_forever_waiting_for_cancellation,
-                    use_shorthand_defined_operation,
+                    op_definition_type,
                 ),
                 start_options=nexusrpc.handler.StartOperationOptions(
                     headers={"my-header-key": "my-header-value"},
                 ),
+                caller_reference=caller_reference,
             ),
             request_cancel,
             task_queue,