8000 add automatic teardown for bootstrappers · modern-python/lite-bootstrap@b59b5a9 · GitHub
[go: up one dir, main page]
Skip to content

Commit b59b5a9

Browse files
lesnik512lesnik512
committed
add automatic teardown for bootstrappers
1 parent 0562a4d commit b59b5a9

File tree

5 files changed

+52
-5
lines changed

5 files changed

+52
-5
lines changed

‎lite_bootstrap/bootstrappers/fastapi_bootstrapper.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import contextlib
12
import dataclasses
23
import typing
34

@@ -20,6 +21,7 @@
2021
if import_checker.is_fastapi_installed:
2122
import fastapi
2223
from fastapi.middleware.cors import CORSMiddleware
24+
from fastapi.routing import _merge_lifespan_context
2325

2426
if import_checker.is_opentelemetry_installed:
2527
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
@@ -160,12 +162,25 @@ class FastAPIBootstrapper(BaseBootstrapper["fastapi.FastAPI"]):
160162
bootstrap_config: FastAPIConfig
161163
not_ready_message = "fastapi is not installed"
162164

165+
@contextlib.asynccontextmanager
166+
async def lifespan_manager(self, _: fastapi.FastAPI) -> typing.AsyncIterator[dict[str, typing.Any]]:
167+
try:
168+
yield {}
169+
finally:
170+
self.teardown()
171+
163172
def __init__(self, bootstrap_config: FastAPIConfig) -> None:
164173
super().__init__(bootstrap_config)
165174
self.bootstrap_config.application.title = bootstrap_config.service_name
166175
self.bootstrap_config.application.debug = bootstrap_config.service_debug
167176
self.bootstrap_config.application.version = bootstrap_config.service_version
168177

178+
old_lifespan_manager = self.bootstrap_config.application.router.lifespan_context
179+
self.bootstrap_config.application.router.lifespan_context = _merge_lifespan_context(
180+
old_lifespan_manager,
181+
self.lifespan_manager,
182+
)
183+
169184
def is_ready(self) -> bool:
170185
return import_checker.is_fastapi_installed
171186

‎lite_bootstrap/bootstrappers/faststream_bootstrapper.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,7 @@ def __init__(self, bootstrap_config: FastStreamConfig) -> None:
156156
super().__init__(bootstrap_config)
157157
if self.bootstrap_config.broker:
158158
self.bootstrap_config.application.broker = self.bootstrap_config.broker
159+
self.bootstrap_config.application.on_shutdown(self.teardown)
159160

160161
def _prepare_application(self) -> "AsgiFastStream":
161162
return self.bootstrap_config.application

‎lite_bootstrap/bootstrappers/litestar_bootstrapper.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -191,6 +191,10 @@ class LitestarBootstrapper(BaseBootstrapper["litestar.Litestar"]):
191191
bootstrap_config: LitestarConfig
192192
not_ready_message = "litestar is not installed"
193193

194+
def __init__(self, bootstrap_config: LitestarConfig) -> None:
195+
super().__init__(bootstrap_config)
196+
self.bootstrap_config.application_config.on_shutdown.append(self.teardown)
197+
194198
def is_ready(self) -> bool:
195199
return import_checker.is_litestar_installed
196200

‎tests/test_fastapi_bootstrap.py

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,9 @@
1+
import asyncio
2+
import contextlib
3+
import dataclasses
4+
import typing
5+
6+
import fastapi
17
import pytest
28
import structlog
39
from opentelemetry.sdk.trace.export import ConsoleSpanExporter
@@ -76,3 +82,21 @@ def test_fastapi_bootstrapper_with_missing_instrument_dependency(
7682
) -> None:
7783
with emulate_package_missing(package_name), pytest.warns(UserWarning, match=package_name):
7884
FastAPIBootstrapper(bootstrap_config=fastapi_config)
85+
86+
87+
def test_fastapi_bootstrap_lifespan(fastapi_config: FastAPIConfig) -> None:
88+
@contextlib.asynccontextmanager
89+
async def lifespan_manager(_: fastapi.FastAPI) -> typing.AsyncIterator[dict[str, typing.Any]]:
90+
try:
91+
yield {}
92+
finally:
93+
await asyncio.sleep(0)
94+
95+
fastapi_config = dataclasses.replace(fastapi_config, application=fastapi.FastAPI(lifespan=lifespan_manager))
96+
bootstrapper = FastAPIBootstrapper(bootstrap_config=fastapi_config)
97+
application = bootstrapper.bootstrap()
98+
99+
with TestClient(application) as test_client:
100+
response = test_client.get(fastapi_config.health_checks_path)
101+
assert response.status_code == status.HTTP_200_OK
102+
assert response.json() == {"health_status": True, "service_name": "microservice", "service_version": "2.0.0"}

‎tests/test_fastapi_offline_docs.py

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -30,14 +30,17 @@ def test_fastapi_offline_docs() -> None:
3030

3131

3232
def test_fastapi_offline_docs_root_path() -> None:
33-
app: FastAPI = FastAPI(title="Tests", root_path="/some-root-path")
33+
app: FastAPI = FastAPI(title="Tests", root_path="/some-root-path", docs_url="/custom_docs")
3434
enable_offline_docs(app)
3535

3636
with TestClient(app, root_path="/some-root-path") as client:
37-
resp = client.get("/docs")
38-
assert resp.status_code == HTTPStatus.OK
39-
assert "/some-root-path/static/swagger-ui.css" in resp.text
40-
assert "/some-root-path/static/swagger-ui-bundle.js" in resp.text
37+
response = client.get("/custom_docs")
38+
assert response.status_code == HTTPStatus.OK
39+
assert "/some-root-path/static/swagger-ui.css" in response.text
40+
assert "/some-root-path/static/swagger-ui-bundle.js" in response.text
41+
42+
response = client.get("/some-root-path/static/swagger-ui.css")
43+
assert response.status_code == HTTPStatus.OK
4144

4245

4346
def test_fastapi_offline_docs_raises_without_openapi_url() -> None:

0 commit comments

Comments
 (0)
0