Rework update_openapi function.

CHANGES.md

@@ -10,6 +10,7 @@
 ### Changed
 
 * Updated CI to test against [pgstac v0.6.12](https://github.com/stac-utils/pgstac/releases/tag/v0.6.12) ([#511](https://github.com/stac-utils/stac-fastapi/pull/511))
+* Reworked `update_openapi` and added a test for it. ([#523](https://github.com/stac-utils/stac-fastapi/pull/523))
 
 ### Removed
 

stac_fastapi/api/stac_fastapi/api/openapi.py

@@ -1,68 +1,36 @@
 """openapi."""
 from fastapi import FastAPI
-from fastapi.openapi.utils import get_openapi
 from starlette.requests import Request
-from starlette.responses import JSONResponse
-
-from stac_fastapi.api.config import ApiExtensions
-from stac_fastapi.types.config import ApiSettings
-
-
-class VndOaiResponse(JSONResponse):
-    """JSON with custom, vendor content-type."""
-
-    media_type = "application/vnd.oai.openapi+json;version=3.0"
+from starlette.responses import JSONResponse, Response
+from starlette.routing import Route, request_response
 
 
 def update_openapi(app: FastAPI) -> FastAPI:
     """Update OpenAPI response content-type.
 
     This function modifies the openapi route to comply with the STAC API spec's
-    required content-type response header
+    required content-type response header.
     """
-    urls = (server_data.get("url") for server_data in app.servers)
-    server_urls = {url for url in urls if url}
-
-    async def openapi(req: Request) -> JSONResponse:
-        root_path = req.scope.get("root_path", "").rstrip("/")
-        if root_path not in server_urls:
-            if root_path and app.root_path_in_servers:
-                app.servers.insert(0, {"url": root_path})
-                server_urls.add(root_path)
-        return VndOaiResponse(app.openapi())
-
-    # Remove the default openapi route
-    app.router.routes = list(
-        filter(lambda r: r.path != app.openapi_url, app.router.routes)
+    # Find the route for the openapi_url in the app.
+    openapi_route: Route = next(
+        route for route in app.router.routes if route.path == app.openapi_url
     )
-    # Add the updated openapi route
-    app.add_route(app.openapi_url, openapi, include_in_schema=False)
-    return app
 
-
-# TODO: Remove or fix, this is currently unused
-# and calls a missing method on ApiSettings
-def config_openapi(app: FastAPI, settings: ApiSettings):
-    """Config openapi."""
-
-    def custom_openapi():
-        """Config openapi."""
-        if app.openapi_schema:
-            return app.openapi_schema
-
-        openapi_schema = get_openapi(
-            title="Arturo STAC API", version="0.1", routes=app.routes
-        )
-
-        if settings.api_extension_is_enabled(ApiExtensions.fields):
-            openapi_schema["paths"]["/search"]["get"]["responses"]["200"]["content"][
-                "application/json"
-            ]["schema"] = {"$ref": "#/components/schemas/ItemCollection"}
-            openapi_schema["paths"]["/search"]["post"]["responses"]["200"]["content"][
-                "application/json"
-            ]["schema"] = {"$ref": "#/components/schemas/ItemCollection"}
-
-        app.openapi_schema = openapi_schema
-        return app.openapi_schema
-
-    app.openapi = custom_openapi
+    # Create a patched endpoint that modifies the content type of the response
+    async def patched_openapi_endpoint(req: Request) -> Response:
+        # Get the response from the original endpoint
+        response: JSONResponse = await openapi_route.endpoint(req)
+        # Update the content type header in place
+        response.headers[
+            "content-type"
+        ] = "application/vnd.oai.openapi+json;version=3.0"
+        # Return the updated response
+        return response
+
+    # When a route is accessed the `handle` function will call `self.app`. So we can
+    # leave the `endpoint` the original function for use in the patched function and
+    # just update the `app` to use the patched function.
+    openapi_route.app = request_response(patched_openapi_endpoint)
+
+    # return the patched app
+    return app

stac_fastapi/api/tests/test_api.py

@@ -49,6 +49,15 @@ def _assert_dependency_applied(api, routes):
                 ), "Authenticated requests should be accepted"
                 assert response.json() == "dummy response"
 
+    def test_openapi_content_type(self):
+        api = self._build_api()
+        with TestClient(api.app) as client:
+            response = client.get(api.settings.openapi_url)
+            assert (
+                response.headers["content-type"]
+                == "application/vnd.oai.openapi+json;version=3.0"
+            )
+
     def test_build_api_with_route_dependencies(self):
         routes = [
             {"path": "/collections", "method": "POST"},