From b59b5a91292f5212efe6bfd4f74653871d1c4f90 Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Tue, 15 Apr 2025 23:19:15 +0300 Subject: [PATCH] add automatic teardown for bootstrappers --- .../bootstrappers/fastapi_bootstrapper.py | 15 ++++++++++++ .../bootstrappers/faststream_bootstrapper.py | 1 + .../bootstrappers/litestar_bootstrapper.py | 4 ++++ tests/test_fastapi_bootstrap.py | 24 +++++++++++++++++++ tests/test_fastapi_offline_docs.py | 13 ++++++---- 5 files changed, 52 insertions(+), 5 deletions(-) diff --git a/lite_bootstrap/bootstrappers/fastapi_bootstrapper.py b/lite_bootstrap/bootstrappers/fastapi_bootstrapper.py index d3ae366..96857cf 100644 --- a/lite_bootstrap/bootstrappers/fastapi_bootstrapper.py +++ b/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 diff --git a/lite_bootstrap/bootstrappers/faststream_bootstrapper.py b/lite_bootstrap/bootstrappers/faststream_bootstrapper.py index 3cce233..49b8107 100644 --- a/lite_bootstrap/bootstrappers/faststream_bootstrapper.py +++ b/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 diff --git a/lite_bootstrap/bootstrappers/litestar_bootstrapper.py b/lite_bootstrap/bootstrappers/litestar_bootstrapper.py index 0810c93..7548f47 100644 --- a/lite_bootstrap/bootstrappers/litestar_bootstrapper.py +++ b/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 diff --git a/tests/test_fastapi_bootstrap.py b/tests/test_fastapi_bootstrap.py index 09e444c..a92413b 100644 --- a/tests/test_fastapi_bootstrap.py +++ b/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"} diff --git a/tests/test_fastapi_offline_docs.py b/tests/test_fastapi_offline_docs.py index 85eceda..37ba2b9 100644 --- a/tests/test_fastapi_offline_docs.py +++ b/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: