Test exception in Nexus start method

Author: dandavison

tests/worker/test_nexus.py

@@ -56,13 +56,15 @@ class OpDefinitionType(IntEnum):
 class SyncResponse:
     op_definition_type: OpDefinitionType
     use_async_def: bool
+    exception_in_operation_start: bool
 
 
 @dataclass
 class AsyncResponse:
     operation_workflow_id: str
     block_forever_waiting_for_cancellation: bool
     op_definition_type: OpDefinitionType
+    exception_in_operation_start: bool
 
 
 # The order of the two types in this union is critical since the data converter matches
@@ -101,6 +103,10 @@ class ServiceInterface:
 #
 
 
+class CustomError(Exception):
+    pass
+
+
 @dataclass
 class HandlerWfInput:
     op_input: OpInput
@@ -136,6 +142,8 @@ async def start(
         OpOutput,
         temporalio.nexus.handler.StartWorkflowOperationResult[HandlerWfOutput],
     ]:
+        if input.response_type.exception_in_operation_start:
+            raise CustomError("Error in Nexus operation start method")
         if isinstance(input.response_type, SyncResponse):
             return OpOutput(
                 value="sync response",
@@ -178,6 +186,8 @@ async def sync_operation(
         self, input: OpInput, options: nexusrpc.handler.StartOperationOptions
     ) -> OpOutput:
         assert isinstance(input.response_type, SyncResponse)
+        if input.response_type.exception_in_operation_start:
+            raise CustomError("Error in Nexus operation start method")
         return OpOutput(
             value="sync response",
             start_options_received_by_handler=options,
@@ -188,6 +198,8 @@ def non_async_sync_operation(
         self, input: OpInput, options: nexusrpc.handler.StartOperationOptions
     ) -> OpOutput:
         assert isinstance(input.response_type, SyncResponse)
+        if input.response_type.exception_in_operation_start:
+            raise CustomError("Error in Nexus operation start method")
         return OpOutput(
             value="sync response",
             start_options_received_by_handler=options,
@@ -198,6 +210,8 @@ async def async_operation(
         self, input: OpInput, options: nexusrpc.handler.StartOperationOptions
     ) -> temporalio.nexus.handler.StartWorkflowOperationResult[HandlerWfOutput]:
         assert isinstance(input.response_type, AsyncResponse)
+        if input.response_type.exception_in_operation_start:
+            raise CustomError("Error in Nexus operation start method")
         return await temporalio.nexus.handler.start_workflow(
             HandlerWorkflow.run,
             args=[HandlerWfInput(op_input=input), options],
@@ -424,6 +438,8 @@ async def run(
 # TODO(dan): nexus endpoint pytest fixture?
 # TODO(dan): get rid of UnsandboxedWorkflowRunner (due to xray)
 # TODO(dan): test headers
+# TODO(dan): enable True for exception_in_operation_start
+@pytest.mark.parametrize("exception_in_operation_start", [False])
 @pytest.mark.parametrize("request_cancel", [False, True])
 @pytest.mark.parametrize(
     "op_definition_type", [OpDefinitionType.SHORTHAND, OpDefinitionType.LONGHAND]
@@ -434,6 +450,7 @@ async def run(
 )
 async def test_sync_response(
     client: Client,
+    exception_in_operation_start: bool,
     request_cancel: bool,
     op_definition_type: OpDefinitionType,
     caller_reference: CallerReference,
@@ -453,7 +470,11 @@ async def test_sync_response(
             args=[
                 CallerWfInput(
                     op_input=OpInput(
-                        response_type=SyncResponse(op_definition_type, True),
+                        response_type=SyncResponse(
+                            op_definition_type=op_definition_type,
+                            use_async_def=True,
+                            exception_in_operation_start=exception_in_operation_start,
+                        ),
                         start_options=nexusrpc.handler.StartOperationOptions(
                             headers={"header-key": "header-value"},
                         ),
@@ -469,11 +490,30 @@ async def test_sync_response(
 
         # The operation result is returned even when request_cancel=True, because the
         # response was synchronous and it could not be cancelled. See explanation below.
-        result = await caller_wf_handle.result()
-        assert result.op_output.value == "sync response"
-        assert result.op_output.start_options_received_by_handler
+        if exception_in_operation_start:
+            with pytest.raises(WorkflowFailureError) as ei:
+                await caller_wf_handle.result()
+            e = ei.value
+            assert isinstance(e, WorkflowFailureError)
+            assert isinstance(e.__cause__, NexusOperationError)
+            assert isinstance(e.__cause__.__cause__, CustomError)
+            # 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 == "ServiceInterface"
+            assert (
+                e.__cause__.operation == "sync_operation"
+                if op_definition_type == OpDefinitionType.SHORTHAND
+                else "sync_or_async_operation"
+            )
+        else:
+            result = await caller_wf_handle.result()
+            assert result.op_output.value == "sync response"
+            assert result.op_output.start_options_received_by_handler
 
 
+# TODO(dan): enable True for exception_in_operation_start
+@pytest.mark.parametrize("exception_in_operation_start", [False])
 @pytest.mark.parametrize("request_cancel", [False, True])
 @pytest.mark.parametrize(
     "op_definition_type", [OpDefinitionType.SHORTHAND, OpDefinitionType.LONGHAND]
@@ -484,6 +524,7 @@ async def test_sync_response(
 )
 async def test_async_response(
     client: Client,
+    exception_in_operation_start: bool,
     request_cancel: bool,
     op_definition_type: OpDefinitionType,
     caller_reference: CallerReference,
@@ -499,7 +540,12 @@ async def test_async_response(
         workflow_failure_exception_types=[Exception],
     ):
         caller_wf_handle, handler_wf_handle = await _start_wf_and_nexus_op(
-            client, task_queue, request_cancel, op_definition_type, caller_reference
+            client,
+            task_queue,
+            exception_in_operation_start,
+            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()
@@ -514,7 +560,23 @@ async def test_async_response(
             caller_wf_handle, handler_wf_handle
         )
 
-        if request_cancel:
+        if exception_in_operation_start:
+            with pytest.raises(WorkflowFailureError) as ei:
+                await caller_wf_handle.result()
+            e = ei.value
+            assert isinstance(e, WorkflowFailureError)
+            assert isinstance(e.__cause__, NexusOperationError)
+            assert isinstance(e.__cause__.__cause__, CustomError)
+            # 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 == "ServiceInterface"
+            assert (
+                e.__cause__.operation == "async_operation"
+                if op_definition_type == OpDefinitionType.SHORTHAND
+                else "sync_or_async_operation"
+            )
+        elif request_cancel:
             # The operation response was asynchronous and so request_cancel is honored. See
             # explanation below.
             with pytest.raises(WorkflowFailureError) as ei:
@@ -549,6 +611,7 @@ async def test_async_response(
 async def _start_wf_and_nexus_op(
     client: Client,
     task_queue: str,
+    exception_in_operation_start: bool,
     request_cancel: bool,
     op_definition_type: OpDefinitionType,
     caller_reference: CallerReference,
@@ -574,6 +637,7 @@ async def _start_wf_and_nexus_op(
                         operation_workflow_id,
                         block_forever_waiting_for_cancellation,
                         op_definition_type,
+                        exception_in_operation_start=exception_in_operation_start,
                     ),
                     start_options=nexusrpc.handler.StartOperationOptions(
                         headers={"header-key": "header-value"},
@@ -605,6 +669,8 @@ async def _start_wf_and_nexus_op(
     return caller_wf_handle, handler_wf_handle
 
 
+# TODO(dan): enable True for exception_in_operation_start
+@pytest.mark.parametrize("exception_in_operation_start", [False])
 @pytest.mark.parametrize(
     "op_definition_type", [OpDefinitionType.SHORTHAND, OpDefinitionType.LONGHAND]
 )
@@ -615,6 +681,7 @@ async def _start_wf_and_nexus_op(
 @pytest.mark.parametrize("response_type", [SyncResponse, AsyncResponse])
 async def test_untyped_caller(
     client: Client,
+    exception_in_operation_start: bool,
     op_definition_type: OpDefinitionType,
     caller_reference: CallerReference,
     response_type: ResponseType,
@@ -629,12 +696,17 @@ async def test_untyped_caller(
         workflow_failure_exception_types=[Exception],
     ):
         if response_type == SyncResponse:
-            response_type = SyncResponse(op_definition_type, True)
+            response_type = SyncResponse(
+                op_definition_type=op_definition_type,
+                use_async_def=True,
+                exception_in_operation_start=exception_in_operation_start,
+            )
         else:
             response_type = AsyncResponse(
-                str(uuid.uuid4()),
-                False,
-                op_definition_type,
+                operation_workflow_id=str(uuid.uuid4()),
+                block_forever_waiting_for_cancellation=False,
+                op_definition_type=op_definition_type,
+                exception_in_operation_start=exception_in_operation_start,
             )
         await create_nexus_endpoint(task_queue, client)
         caller_wf_handle = await client.start_workflow(