From f369e83c5301a8874250d0e55d32f36953c43002 Mon Sep 17 00:00:00 2001 From: Thomas Pelletier Date: Wed, 31 Jan 2024 09:15:37 -0500 Subject: [PATCH] Remove mount point and attach to root --- src/dispatch/fastapi.py | 22 +++++++++------------- tests/test_fastapi.py | 19 ++++++++----------- 2 files changed, 17 insertions(+), 24 deletions(-) diff --git a/src/dispatch/fastapi.py b/src/dispatch/fastapi.py index 4e0de099..38d4aaa1 100644 --- a/src/dispatch/fastapi.py +++ b/src/dispatch/fastapi.py @@ -30,19 +30,16 @@ def read_root(): def configure( app: fastapi.FastAPI, api_key: None | str = None, - mount_path: str = "/dispatch", ): """Configure the FastAPI app to use Dispatch programmable endpoints. - It mounts a sub-app at the given mount path that implements the Dispatch - interface. It also adds a a decorator named @app.dispatch_coroutine() to - register coroutines. + It mounts a sub-app that implements the Dispatch gRPC interface. It also + adds a a decorator named @app.dispatch_coroutine() to register coroutines. Args: app: The FastAPI app to configure. api_key: Dispatch API key to use for authentication. Uses the value of the DISPATCH_API_KEY environment variable by default. - mount_path: The path to mount Dispatch programmable endpoints at. Raises: ValueError: If any of the required arguments are missing. @@ -53,13 +50,11 @@ def configure( raise ValueError("app is required") if not api_key: raise ValueError("api_key is required") - if not mount_path: - raise ValueError("mount_path is required") dispatch_app = _new_app() app.__setattr__("dispatch_coroutine", dispatch_app.dispatch_coroutine) - app.mount(mount_path, dispatch_app) + app.mount("/ring.coroutine.v1.ExecutorService", dispatch_app) class _DispatchAPI(fastapi.FastAPI): @@ -96,15 +91,11 @@ def _new_app(): app = _DispatchAPI() app._coroutines = {} - @app.get("/", response_class=fastapi.responses.PlainTextResponse) - def read_root(): - return "ok" - @app.post( # The endpoint for execution is hardcoded at the moment. If the service # gains more endpoints, this should be turned into a dynamic dispatch # like the official gRPC server does. - "/ring.coroutine.v1.ExecutorService/Execute", + "/Execute", response_class=_GRPCResponse, ) async def execute(request: fastapi.Request): @@ -115,6 +106,11 @@ async def execute(request: fastapi.Request): req = ring.coroutine.v1.coroutine_pb2.ExecuteRequest.FromString(data) + if not req.coroutine_uri: + raise fastapi.HTTPException( + status_code=400, detail="coroutine_uri is required" + ) + # TODO: be more graceful. This will crash if the coroutine is not found, # and the coroutine version is not taken into account. coroutine = app._coroutines[ diff --git a/tests/test_fastapi.py b/tests/test_fastapi.py index d7c16a72..ec48a790 100644 --- a/tests/test_fastapi.py +++ b/tests/test_fastapi.py @@ -26,11 +26,11 @@ def read_root(): # Ensure existing routes are still working. resp = client.get("/") self.assertEqual(resp.status_code, 200) + self.assertEqual(resp.json(), {"Hello": "World"}) - # Ensure Dispatch root is working. - resp = client.get("/dispatch/") - self.assertEqual(resp.status_code, 200) - self.assertEqual(resp.text, "ok") + # Ensure Dispatch root is available. + resp = client.post("/ring.coroutine.v1.ExecutorService/Execute") + self.assertEqual(resp.status_code, 400) def test_configure_no_app(self): with self.assertRaises(ValueError): @@ -41,13 +41,9 @@ def test_configure_no_api_key(self): with self.assertRaises(ValueError): dispatch.fastapi.configure(app, api_key=None) - def test_configure_no_mount_path(self): - app = fastapi.FastAPI() - with self.assertRaises(ValueError): - dispatch.fastapi.configure(app, api_key="test-key", mount_path=None) - def test_fastapi_simple_request(self): - app = dispatch.fastapi._new_app() + app = fastapi.FastAPI() + dispatch.fastapi.configure(app, api_key="test-key") @app.dispatch_coroutine() def my_cool_coroutine(input: Input) -> Output: @@ -89,7 +85,8 @@ def response_output(resp: coroutine_pb2.ExecuteResponse) -> Any: class TestCoroutine(unittest.TestCase): def setUp(self): - self.app = dispatch.fastapi._new_app() + self.app = fastapi.FastAPI() + dispatch.fastapi.configure(self.app, api_key="test-key") http_client = TestClient(self.app) self.client = executor_service.client(http_client)