Fix status serialization

dispatchrun · Jan 30, 2024 · feb7001 · feb7001
1 parent c5efaf7
commit feb7001
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/src/dispatch/coroutine.py b/src/dispatch/coroutine.py
@@ -203,9 +203,10 @@ def error(cls, error: Error) -> Output:
         """Terminally exit the coroutine with the provided error."""
         return Output(
             coroutine_pb2.ExecuteResponse(
+                status=error.status._proto,
                 exit=coroutine_pb2.Exit(
                     result=coroutine_pb2.Result(error=error._as_proto())
-                )
+                ),
             )
         )
 
@@ -298,6 +299,7 @@ def __init__(self, status: Status, type: str | None, message: str | None):
 
         self.type = type
         self.message = message
+        self.status = status
 
     @classmethod
     def from_exception(cls, ex: Exception, status: Status | None = None) -> Error:

diff --git a/tests/test_fastapi.py b/tests/test_fastapi.py
@@ -269,6 +269,7 @@ def mycoro(input: Input) -> Output:
         resp = self.execute(mycoro)
         self.assertEqual("ZeroDivisionError", resp.exit.result.error.type)
         self.assertEqual("division by zero", resp.exit.result.error.message)
+        self.assertEqual(Status.TEMPORARY_ERROR, resp.status)
 
     def test_coroutine_unexpected_exception(self):
         @self.app.dispatch_coroutine()
@@ -279,3 +280,14 @@ def mycoro(input: Input) -> Output:
         resp = self.execute(mycoro)
         self.assertEqual("ZeroDivisionError", resp.exit.result.error.type)
         self.assertEqual("division by zero", resp.exit.result.error.message)
+        self.assertEqual(Status.TEMPORARY_ERROR, resp.status)
+
+    def test_specific_status(self):
+        @self.app.dispatch_coroutine()
+        def mycoro(input: Input) -> Output:
+            return Output.error(Error(Status.THROTTLED, "foo", "bar"))
+
+        resp = self.execute(mycoro)
+        self.assertEqual("foo", resp.exit.result.error.type)
+        self.assertEqual("bar", resp.exit.result.error.message)
+        self.assertEqual(Status.THROTTLED, resp.status)