add automatic teardown for bootstrappers

lite_bootstrap/bootstrappers/fastapi_bootstrapper.py

@@ -1,3 +1,4 @@
+import contextlib
 import dataclasses
 import typing
 
@@ -20,6 +21,7 @@
 if import_checker.is_fastapi_installed:
     import fastapi
     from fastapi.middleware.cors import CORSMiddleware
+    from fastapi.routing import _merge_lifespan_context
 
 if import_checker.is_opentelemetry_installed:
     from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
@@ -160,12 +162,25 @@ class FastAPIBootstrapper(BaseBootstrapper["fastapi.FastAPI"]):
     bootstrap_config: FastAPIConfig
     not_ready_message = "fastapi is not installed"
 
+    @contextlib.asynccontextmanager
+    async def lifespan_manager(self, _: fastapi.FastAPI) -> typing.AsyncIterator[dict[str, typing.Any]]:
+        try:
+            yield {}
+        finally:
+            self.teardown()
+
     def __init__(self, bootstrap_config: FastAPIConfig) -> None:
         super().__init__(bootstrap_config)
         self.bootstrap_config.application.title = bootstrap_config.service_name
         self.bootstrap_config.application.debug = bootstrap_config.service_debug
         self.bootstrap_config.application.version = bootstrap_config.service_version
 
+        old_lifespan_manager = self.bootstrap_config.application.router.lifespan_context
+        self.bootstrap_config.application.router.lifespan_context = _merge_lifespan_context(
+            old_lifespan_manager,
+            self.lifespan_manager,
+        )
+
     def is_ready(self) -> bool:
         return import_checker.is_fastapi_installed
 

lite_bootstrap/bootstrappers/faststream_bootstrapper.py

@@ -156,6 +156,7 @@ def __init__(self, bootstrap_config: FastStreamConfig) -> None:
         super().__init__(bootstrap_config)
         if self.bootstrap_config.broker:
             self.bootstrap_config.application.broker = self.bootstrap_config.broker
+        self.bootstrap_config.application.on_shutdown(self.teardown)
 
     def _prepare_application(self) -> "AsgiFastStream":
         return self.bootstrap_config.application

lite_bootstrap/bootstrappers/litestar_bootstrapper.py

@@ -191,6 +191,10 @@ class LitestarBootstrapper(BaseBootstrapper["litestar.Litestar"]):
     bootstrap_config: LitestarConfig
     not_ready_message = "litestar is not installed"
 
+    def __init__(self, bootstrap_config: LitestarConfig) -> None:
+        super().__init__(bootstrap_config)
+        self.bootstrap_config.application_config.on_shutdown.append(self.teardown)
+
     def is_ready(self) -> bool:
         return import_checker.is_litestar_installed
 

tests/test_fastapi_bootstrap.py

@@ -1,3 +1,9 @@
+import asyncio
+import contextlib
+import dataclasses
+import typing
+
+import fastapi
 import pytest
 import structlog
 from opentelemetry.sdk.trace.export import ConsoleSpanExporter
@@ -76,3 +82,21 @@ def test_fastapi_bootstrapper_with_missing_instrument_dependency(
 ) -> None:
     with emulate_package_missing(package_name), pytest.warns(UserWarning, match=package_name):
         FastAPIBootstrapper(bootstrap_config=fastapi_config)
+
+
+def test_fastapi_bootstrap_lifespan(fastapi_config: FastAPIConfig) -> None:
+    @contextlib.asynccontextmanager
+    async def lifespan_manager(_: fastapi.FastAPI) -> typing.AsyncIterator[dict[str, typing.Any]]:
+        try:
+            yield {}
+        finally:
+            await asyncio.sleep(0)
+
+    fastapi_config = dataclasses.replace(fastapi_config, application=fastapi.FastAPI(lifespan=lifespan_manager))
+    bootstrapper = FastAPIBootstrapper(bootstrap_config=fastapi_config)
+    application = bootstrapper.bootstrap()
+
+    with TestClient(application) as test_client:
+        response = test_client.get(fastapi_config.health_checks_path)
+        assert response.status_code == status.HTTP_200_OK
+        assert response.json() == {"health_status": True, "service_name": "microservice", "service_version": "2.0.0"}

tests/test_fastapi_offline_docs.py

@@ -30,14 +30,17 @@ def test_fastapi_offline_docs() -> None:
 
 
 def test_fastapi_offline_docs_root_path() -> None:
-    app: FastAPI = FastAPI(title="Tests", root_path="/some-root-path")
+    app: FastAPI = FastAPI(title="Tests", root_path="/some-root-path", docs_url="/custom_docs")
     enable_offline_docs(app)
 
     with TestClient(app, root_path="/some-root-path") as client:
-        resp = client.get("/docs")
-        assert resp.status_code == HTTPStatus.OK
-        assert "/some-root-path/static/swagger-ui.css" in resp.text
-        assert "/some-root-path/static/swagger-ui-bundle.js" in resp.text
+        response = client.get("/custom_docs")
+        assert response.status_code == HTTPStatus.OK
+        assert "/some-root-path/static/swagger-ui.css" in response.text
+        assert "/some-root-path/static/swagger-ui-bundle.js" in response.text
+
+        response = client.get("/some-root-path/static/swagger-ui.css")
+        assert response.status_code == HTTPStatus.OK
 
 
 def test_fastapi_offline_docs_raises_without_openapi_url() -> None: