refactor: expose reusable lifespan handler

Author: alukach

src/stac_auth_proxy/__init__.py

@@ -8,12 +8,11 @@
 
 from .app import configure_app, create_app
 from .config import Settings
-from .lifespan import check_conformance, check_server_health
+from .lifespan import lifespan
 
 __all__ = [
     "create_app",
     "configure_app",
-    "check_conformance",
-    "check_server_health",
+    "lifespan",
     "Settings",
 ]

src/stac_auth_proxy/app.py

@@ -6,7 +6,6 @@
 """
 
 import logging
-from contextlib import asynccontextmanager
 from typing import Optional
 
 from fastapi import FastAPI
@@ -26,7 +25,7 @@
     ProcessLinksMiddleware,
     RemoveRootPathMiddleware,
 )
-from .utils.lifespan import check_conformance, check_server_health
+from .lifespan import lifespan
 
 logger = logging.getLogger(__name__)
 
@@ -136,48 +135,13 @@ def configure_app(app: FastAPI, settings: Optional[Settings] = None) -> FastAPI:
         )
 
     return app
-
-
 def create_app(settings: Optional[Settings] = None) -> FastAPI:
     """FastAPI Application Factory."""
     settings = settings or Settings()
 
-    #
-    # Application
-    #
-
-    @asynccontextmanager
-    async def lifespan(app: FastAPI):
-        assert settings
-
-        # Wait for upstream servers to become available
-        if settings.wait_for_upstream:
-            logger.info("Running upstream server health checks...")
-            urls = [settings.upstream_url, settings.oidc_discovery_internal_url]
-            for url in urls:
-                await check_server_health(url=url)
-            logger.info(
-                "Upstream servers are healthy:\n%s",
-                "\n".join([f" - {url}" for url in urls]),
-            )
-
-        # Log all middleware connected to the app
-        logger.info(
-            "Connected middleware:\n%s",
-            "\n".join([f" - {m.cls.__name__}" for m in app.user_middleware]),
-        )
-
-        if settings.check_conformance:
-            await check_conformance(
-                app.user_middleware,
-                str(settings.upstream_url),
-            )
-
-        yield
-
     app = FastAPI(
         openapi_url=None,  # Disable OpenAPI schema endpoint, we want to serve upstream's schema
-        lifespan=lifespan,
+        lifespan=lifespan(settings),
         root_path=settings.root_path,
     )
     if app.root_path:

src/stac_auth_proxy/lifespan.py

@@ -1,10 +1,58 @@
-"""Public access to lifespan health checks.
+"""Reusable lifespan handler for FastAPI applications."""
 
-This module re-exports the ``check_server_health`` and ``check_conformance``
-utilities so that library users can import them without reaching into the
-internal ``utils`` package.
-"""
+from contextlib import asynccontextmanager
+import logging
+from fastapi import FastAPI
 
+from .config import Settings
 from .utils.lifespan import check_conformance, check_server_health
 
-__all__ = ["check_server_health", "check_conformance"]
+logger = logging.getLogger(__name__)
+
+
+def lifespan(settings: Settings | None = None):
+    """Create a lifespan handler that runs startup checks.
+
+    Parameters
+    ----------
+    settings : Settings | None
+        Configuration for health and conformance checks. If omitted, default
+        settings are loaded.
+
+    Returns
+    -------
+    Callable[[FastAPI], AsyncContextManager[Any]]
+        A callable suitable for the ``lifespan`` parameter of ``FastAPI``.
+    """
+
+    settings = settings or Settings()
+
+    @asynccontextmanager
+    async def _lifespan(app: FastAPI):
+        # Wait for upstream servers to become available
+        if settings.wait_for_upstream:
+            logger.info("Running upstream server health checks...")
+            urls = [settings.upstream_url, settings.oidc_discovery_internal_url]
+            for url in urls:
+                await check_server_health(url=url)
+            logger.info(
+                "Upstream servers are healthy:\n%s",
+                "\n".join([f" - {url}" for url in urls]),
+            )
+
+        # Log all middleware connected to the app
+        logger.info(
+            "Connected middleware:\n%s",
+            "\n".join([f" - {m.cls.__name__}" for m in app.user_middleware]),
+        )
+
+        if settings.check_conformance:
+            await check_conformance(app.user_middleware, str(settings.upstream_url))
+
+        yield
+
+    return _lifespan
+
+
+__all__ = ["lifespan"]
+

tests/test_lifespan.py

@@ -1,13 +1,16 @@
 """Tests for lifespan module."""
 
 from dataclasses import dataclass
-from unittest.mock import patch
+from unittest.mock import AsyncMock, patch
 
 import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
 from starlette.middleware import Middleware
 from starlette.types import ASGIApp
 
-from stac_auth_proxy import check_conformance, check_server_health
+from stac_auth_proxy import Settings, lifespan as lifespan_handler
+from stac_auth_proxy.utils.lifespan import check_conformance, check_server_health
 from stac_auth_proxy.utils.middleware import required_conformance
 
 
@@ -80,3 +83,25 @@ def __init__(self, app):
 
     middleware = [Middleware(NoConformanceMiddleware)]
     await check_conformance(middleware, source_api_server)
+
+
+def test_lifespan_reusable():
+    """Ensure the public lifespan handler runs health and conformance checks."""
+    settings = Settings(
+        upstream_url="https://example.com",
+        oidc_discovery_url="https://example.com/.well-known/openid-configuration",
+    )
+    with patch(
+        "stac_auth_proxy.lifespan.check_server_health",
+        new=AsyncMock(),
+    ) as mock_health, patch(
+        "stac_auth_proxy.lifespan.check_conformance",
+        new=AsyncMock(),
+    ) as mock_conf:
+        app = FastAPI(lifespan=lifespan_handler(settings))
+        with TestClient(app):
+            pass
+        assert mock_health.await_count == 2
+        mock_conf.assert_awaited_once_with(
+            app.user_middleware, str(settings.upstream_url)
+        )