Skip to content

Commit b81a6fb

Browse files
pcrespovpcrespov
committed
tests concept further
1 parent 11a7e7b commit b81a6fb

File tree

1 file changed

+221
-3
lines changed

1 file changed

+221
-3
lines changed

‎packages/service-library/tests/fastapi/test_lifespan_utils.py‎

Lines changed: 221 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,21 @@
1+
# pylint: disable=protected-access
2+
# pylint: disable=redefined-outer-name
3+
# pylint: disable=too-many-arguments
4+
# pylint: disable=unused-argument
5+
# pylint: disable=unused-variable
6+
7+
8+
import logging
19
from collections.abc import AsyncIterator
10+
from typing import Any
211

3-
import asgi_lifespan
412
import pytest
13+
from asgi_lifespan import LifespanManager as ASGILifespanManager
14+
from common_library.errors_classes import OsparcErrorMixin
515
from fastapi import FastAPI
6-
from fastapi_lifespan_manager import State
16+
from fastapi_lifespan_manager import LifespanManager, State
17+
from pytest_mock import MockerFixture
18+
from pytest_simcore.helpers.logging_tools import log_context
719
from servicelib.fastapi.lifespan_utils import combine_lifespans
820

921

@@ -24,7 +36,7 @@ async def cache_lifespan(app: FastAPI) -> AsyncIterator[State]:
2436

2537
capsys.readouterr()
2638

27-
async with asgi_lifespan.LifespanManager(app):
39+
async with ASGILifespanManager(app):
2840
messages = capsys.readouterr().out
2941

3042
assert "setup DB" in messages
@@ -38,3 +50,209 @@ async def cache_lifespan(app: FastAPI) -> AsyncIterator[State]:
3850
assert "setup CACHE" not in messages
3951
assert "shutdown DB" in messages
4052
assert "shutdown CACHE" in messages
53+
54+
55+
@pytest.fixture
56+
def postgres_lifespan() -> LifespanManager:
57+
lifespan_manager = LifespanManager()
58+
59+
@lifespan_manager.add
60+
async def _setup_postgres_sync_engine(_) -> AsyncIterator[State]:
61+
with log_context(logging.INFO, "postgres_sync_engine"):
62+
# pass state to children
63+
yield {"postgres": {"engine": "Some Engine"}}
64+
65+
@lifespan_manager.add
66+
async def _setup_postgres_async_engine(_, state: State) -> AsyncIterator[State]:
67+
with log_context(logging.INFO, "postgres_async_engine"):
68+
# pass state to children
69+
70+
state["postgres"].update(aengine="Some Async Engine")
71+
yield state
72+
73+
return lifespan_manager
74+
75+
76+
@pytest.fixture
77+
def rabbitmq_lifespan() -> LifespanManager:
78+
lifespan_manager = LifespanManager()
79+
80+
@lifespan_manager.add
81+
async def _setup_rabbitmq(app: FastAPI) -> AsyncIterator[State]:
82+
with log_context(logging.INFO, "rabbitmq"):
83+
84+
with pytest.raises(AttributeError, match="rabbitmq_rpc_server"):
85+
_ = app.state.rabbitmq_rpc_server
86+
87+
# pass state to children
88+
yield {"rabbitmq_rpc_server": "Some RabbitMQ RPC Server"}
89+
90+
return lifespan_manager
91+
92+
93+
async def test_app_lifespan_composition(
94+
postgres_lifespan: LifespanManager, rabbitmq_lifespan: LifespanManager
95+
):
96+
# The app has its own database and rpc-server to initialize
97+
# this is how you connect the lifespans pre-defined in servicelib
98+
99+
@postgres_lifespan.add
100+
async def setup_database(app: FastAPI, state: State) -> AsyncIterator[State]:
101+
102+
with log_context(logging.INFO, "app database"):
103+
assert state["postgres"] == {
104+
"engine": "Some Engine",
105+
"aengine": "Some Async Engine",
106+
}
107+
108+
with pytest.raises(AttributeError, match="database_engine"):
109+
_ = app.state.database_engine
110+
111+
app.state.database_engine = state["postgres"]["engine"]
112+
113+
yield {}
114+
115+
# tear-down stage
116+
assert app.state.database_engine
117+
118+
@rabbitmq_lifespan.add
119+
async def setup_rpc_server(app: FastAPI, state: State) -> AsyncIterator[State]:
120+
with log_context(logging.INFO, "app rpc-server"):
121+
assert "rabbitmq_rpc_server" in state
122+
123+
app.state.rpc_server = state["rabbitmq_rpc_server"]
124+
125+
yield {}
126+
127+
# Composes lifepans
128+
app_lifespan = LifespanManager()
129+
app_lifespan.include(postgres_lifespan)
130+
app_lifespan.include(rabbitmq_lifespan)
131+
132+
app = FastAPI(lifespan=app_lifespan)
133+
async with ASGILifespanManager(app) as asgi_manager:
134+
135+
# asgi_manage state
136+
assert asgi_manager._state == { # noqa: SLF001
137+
"postgres": {
138+
"engine": "Some Engine",
139+
"aengine": "Some Async Engine",
140+
},
141+
"rabbitmq_rpc_server": "Some RabbitMQ RPC Server",
142+
}
143+
144+
# app state
145+
assert app.state.database_engine
146+
assert app.state.rpc_server
147+
148+
# Logs shows lifespan execution:
149+
# -> postgres_sync_engine starting ...
150+
# -> postgres_async_engine starting ...
151+
# -> app database starting ...
152+
# -> rabbitmq starting ...
153+
# -> app rpc-server starting ...
154+
# <- app rpc-server done ()
155+
# <- rabbitmq done ()
156+
# <- app database done (1ms)
157+
# <- postgres_async_engine done (1ms)
158+
# <- postgres_sync_engine done (1ms)
159+
160+
161+
class LifespanError(OsparcErrorMixin, RuntimeError): ...
162+
163+
164+
class LifespanOnStartupError(LifespanError):
165+
msg_template = "Failed during startup of {module}"
166+
167+
168+
class LifespanOnShutdownError(LifespanError):
169+
msg_template = "Failed during shutdown of {module}"
170+
171+
172+
@pytest.fixture
173+
def failing_lifespan_manager(mocker: MockerFixture):
174+
startup_step = mocker.MagicMock()
175+
shutdown_step = mocker.MagicMock()
176+
handle_error = mocker.MagicMock()
177+
178+
def raise_error():
179+
msg = "failing module"
180+
raise RuntimeError(msg)
181+
182+
async def setup_failing_on_startup(app: FastAPI) -> AsyncIterator[State]:
183+
_name = setup_failing_on_startup.__name__
184+
185+
with log_context(logging.INFO, _name):
186+
try:
187+
raise_error()
188+
startup_step(_name)
189+
except RuntimeError as exc:
190+
handle_error(_name, exc)
191+
raise LifespanOnStartupError(module=_name) from exc
192+
yield {}
193+
shutdown_step(_name)
194+
195+
async def setup_failing_on_shutdown(app: FastAPI) -> AsyncIterator[State]:
196+
_name = setup_failing_on_shutdown.__name__
197+
198+
with log_context(logging.INFO, _name):
199+
startup_step(_name)
200+
yield {}
201+
try:
202+
raise_error()
203+
shutdown_step(_name)
204+
except RuntimeError as exc:
205+
handle_error(_name, exc)
206+
raise LifespanOnShutdownError(module=_name) from exc
207+
208+
return {
209+
"startup_step": startup_step,
210+
"shutdown_step": shutdown_step,
211+
"handle_error": handle_error,
212+
"setup_failing_on_startup": setup_failing_on_startup,
213+
"setup_failing_on_shutdown": setup_failing_on_shutdown,
214+
}
215+
216+
217+
async def test_app_lifespan_with_error_on_startup(
218+
failing_lifespan_manager: dict[str, Any],
219+
):
220+
app_lifespan = LifespanManager()
221+
app_lifespan.add(failing_lifespan_manager["setup_failing_on_startup"])
222+
app = FastAPI(lifespan=app_lifespan)
223+
224+
with pytest.raises(LifespanOnStartupError) as err_info:
225+
async with ASGILifespanManager(app):
226+
...
227+
228+
exception = err_info.value
229+
assert failing_lifespan_manager["handle_error"].called
230+
assert not failing_lifespan_manager["startup_step"].called
231+
assert not failing_lifespan_manager["shutdown_step"].called
232+
assert exception.error_context() == {
233+
"module": "setup_failing_on_startup",
234+
"message": "Failed during startup of setup_failing_on_startup",
235+
"code": "RuntimeError.LifespanError.LifespanOnStartupError",
236+
}
237+
238+
239+
async def test_app_lifespan_with_error_on_shutdown(
240+
failing_lifespan_manager: dict[str, Any],
241+
):
242+
app_lifespan = LifespanManager()
243+
app_lifespan.add(failing_lifespan_manager["setup_failing_on_shutdown"])
244+
app = FastAPI(lifespan=app_lifespan)
245+
246+
with pytest.raises(LifespanOnShutdownError) as err_info:
247+
async with ASGILifespanManager(app):
248+
...
249+
250+
exception = err_info.value
251+
assert failing_lifespan_manager["handle_error"].called
252+
assert failing_lifespan_manager["startup_step"].called
253+
assert not failing_lifespan_manager["shutdown_step"].called
254+
assert exception.error_context() == {
255+
"module": "setup_failing_on_shutdown",
256+
"message": "Failed during shutdown of setup_failing_on_shutdown",
257+
"code": "RuntimeError.LifespanError.LifespanOnShutdownError",
258+
}

0 commit comments

Comments
 (0)