✨ Refactor lifespan utility functions and enhance error handling in FastAPI

Author: pcrespov

packages/service-library/src/servicelib/fastapi/lifespan_utils.py

@@ -20,25 +20,37 @@ class LifespanAlreadyCalledError(LifespanError):
     msg_template = "The lifespan '{lifespan_name}' has already been called."
 
 
+class LifespanExpectedCalledError(LifespanError):
+    msg_template = "The lifespan '{lifespan_name}' was not called. Ensure it is properly configured and invoked."
+
+
 _CALLED_LIFESPANS_KEY: Final[str] = "_CALLED_LIFESPANS"
 
 
 def is_lifespan_called(state: State, lifespan_name: str) -> bool:
+    assert not isinstance(  # nosec
+        state, FastAPI
+    ), "TIP: lifespan func has (app, state) positional arguments"
+
     called_lifespans = state.get(_CALLED_LIFESPANS_KEY, set())
     return lifespan_name in called_lifespans
 
 
-def record_lifespan_called_once(state: State, lifespan_name: str) -> State:
+def mark_lifespace_called(state: State, lifespan_name: str) -> State:
     """Validates if a lifespan has already been called and records it in the state.
     Raises LifespanAlreadyCalledError if the lifespan has already been called.
     """
-    assert not isinstance(  # nosec
-        state, FastAPI
-    ), "TIP: lifespan func has (app, state) positional arguments"
-
     if is_lifespan_called(state, lifespan_name):
         raise LifespanAlreadyCalledError(lifespan_name=lifespan_name)
 
     called_lifespans = state.get(_CALLED_LIFESPANS_KEY, set())
     called_lifespans.add(lifespan_name)
     return {_CALLED_LIFESPANS_KEY: called_lifespans}
+
+
+def ensure_lifespan_called(state: State, lifespan_name: str) -> None:
+    """Ensures that a lifespan has been called.
+    Raises LifespanNotCalledError if the lifespan has not been called.
+    """
+    if not is_lifespan_called(state, lifespan_name):
+        raise LifespanExpectedCalledError(lifespan_name=lifespan_name)

packages/service-library/src/servicelib/fastapi/postgres_lifespan.py

@@ -10,7 +10,7 @@
 from sqlalchemy.ext.asyncio import AsyncEngine
 
 from ..db_asyncpg_utils import create_async_engine_and_database_ready
-from .lifespan_utils import LifespanOnStartupError, record_lifespan_called_once
+from .lifespan_utils import LifespanOnStartupError, mark_lifespace_called
 
 _logger = logging.getLogger(__name__)
 
@@ -33,7 +33,7 @@ async def postgres_database_lifespan(_: FastAPI, state: State) -> AsyncIterator[
     with log_context(_logger, logging.INFO, f"{__name__}"):
 
         # Mark lifespan as called
-        called_state = record_lifespan_called_once(state, "postgres_database_lifespan")
+        called_state = mark_lifespace_called(state, "postgres_database_lifespan")
 
         settings = state[PostgresLifespanState.POSTGRES_SETTINGS]
 

packages/service-library/src/servicelib/fastapi/rabbitmq_lifespan.py

@@ -8,7 +8,10 @@
 from servicelib.rabbitmq import wait_till_rabbitmq_responsive
 from settings_library.rabbit import RabbitSettings
 
-from .lifespan_utils import LifespanOnStartupError, record_lifespan_called_once
+from .lifespan_utils import (
+    LifespanOnStartupError,
+    mark_lifespace_called,
+)
 
 _logger = logging.getLogger(__name__)
 
@@ -33,7 +36,7 @@ async def rabbitmq_connectivity_lifespan(
     with log_context(_logger, logging.INFO, _lifespan_name):
 
         # Check if lifespan has already been called
-        called_state = record_lifespan_called_once(state, _lifespan_name)
+        called_state = mark_lifespace_called(state, _lifespan_name)
 
         # Validate input state
         try:

packages/service-library/src/servicelib/fastapi/redis_lifespan.py

@@ -10,7 +10,7 @@
 from settings_library.redis import RedisDatabase, RedisSettings
 
 from ..redis import RedisClientSDK
-from .lifespan_utils import LifespanOnStartupError, record_lifespan_called_once
+from .lifespan_utils import LifespanOnStartupError, mark_lifespace_called
 
 _logger = logging.getLogger(__name__)
 
@@ -34,7 +34,7 @@ async def redis_database_lifespan(_: FastAPI, state: State) -> AsyncIterator[Sta
     with log_context(_logger, logging.INFO, f"{__name__}"):
 
         # Check if lifespan has already been called
-        called_state = record_lifespan_called_once(state, "redis_database_lifespan")
+        called_state = mark_lifespace_called(state, "redis_database_lifespan")
 
         # Validate input state
         try:

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

@@ -17,9 +17,11 @@
 from pytest_simcore.helpers.logging_tools import log_context
 from servicelib.fastapi.lifespan_utils import (
     LifespanAlreadyCalledError,
+    LifespanExpectedCalledError,
     LifespanOnShutdownError,
     LifespanOnStartupError,
-    record_lifespan_called_once,
+    ensure_lifespan_called,
+    mark_lifespace_called,
 )
 
 
@@ -259,18 +261,21 @@ async def test_app_lifespan_with_error_on_shutdown(
 
 
 async def test_lifespan_called_more_than_once(is_pdb_enabled: bool):
-    state = {}
-
     app_lifespan = LifespanManager()
 
     @app_lifespan.add
     async def _one(_, state: State) -> AsyncIterator[State]:
-        called_state = record_lifespan_called_once(state, "test_lifespan_one")
+        called_state = mark_lifespace_called(state, "test_lifespan_one")
         yield {"other": 0, **called_state}
 
     @app_lifespan.add
     async def _two(_, state: State) -> AsyncIterator[State]:
-        called_state = record_lifespan_called_once(state, "test_lifespan_two")
+        ensure_lifespan_called(state, "test_lifespan_one")
+
+        with pytest.raises(LifespanExpectedCalledError):
+            ensure_lifespan_called(state, "test_lifespan_three")
+
+        called_state = mark_lifespace_called(state, "test_lifespan_two")
         yield {"something": 0, **called_state}
 
     app_lifespan.add(_one)  # added "by mistake"

packages/service-library/tests/fastapi/test_rabbitmq_lifespan.py

@@ -37,11 +37,14 @@ def mock_rabbitmq_connection(mocker: MockerFixture) -> MockType:
 
 @pytest.fixture
 def mock_rabbitmq_rpc_client_class(mocker: MockerFixture) -> MockType:
-    return mocker.patch.object(
-        servicelib.rabbitmq._client_rpc,
-        "RabbitMQRPCClient",
-        return_value=mocker.AsyncMock(),
+    mock_rpc_client_instance = mocker.AsyncMock()
+    mocker.patch.object(
+        servicelib.rabbitmq._client_rpc.RabbitMQRPCClient,
+        "create",
+        return_value=mock_rpc_client_instance,
     )
+    mock_rpc_client_instance.close = mocker.AsyncMock()
+    return mock_rpc_client_instance
 
 
 @pytest.fixture
@@ -64,25 +67,40 @@ class AppSettings(BaseApplicationSettings):
             ..., json_schema_extra={"auto_default_from_env": True}
         )
 
+    # setup settings
     async def my_app_settings(app: FastAPI) -> AsyncIterator[State]:
         app.state.settings = AppSettings.create_from_envs()
 
         yield RabbitMQLifespanState(
             RABBIT_SETTINGS=app.state.settings.RABBITMQ,
         ).model_dump()
 
+    # setup rpc-server using rabbitmq_rpc_client_context (yes, a "rpc_server" is built with an RabbitMQRpcClient)
     async def my_app_rpc_server(app: FastAPI, state: State) -> AsyncIterator[State]:
+        assert "RABBIT_CONNECTIVITY_LIFESPAN_NAME" in state
 
         async with rabbitmq_rpc_client_context(
             "rpc_server", app.state.settings.RABBITMQ
         ) as rpc_server:
             app.state.rpc_server = rpc_server
             yield {}
 
+    # setup rpc-client using rabbitmq_rpc_client_context
+    async def my_app_rpc_client(app: FastAPI, state: State) -> AsyncIterator[State]:
+
+        assert "RABBIT_CONNECTIVITY_LIFESPAN_NAME" in state
+
+        async with rabbitmq_rpc_client_context(
+            "rpc_client", app.state.settings.RABBITMQ
+        ) as rpc_client:
+            app.state.rpc_client = rpc_client
+            yield {}
+
     app_lifespan = LifespanManager()
     app_lifespan.add(my_app_settings)
     app_lifespan.add(rabbitmq_connectivity_lifespan)
     app_lifespan.add(my_app_rpc_server)
+    app_lifespan.add(my_app_rpc_client)
 
     assert not mock_rabbitmq_connection.called
     assert not mock_rabbitmq_rpc_client_class.called
@@ -103,17 +121,20 @@ async def test_lifespan_rabbitmq_in_an_app(
         app,
         startup_timeout=None if is_pdb_enabled else 10,
         shutdown_timeout=None if is_pdb_enabled else 10,
-    ) as asgi_manager:
+    ):
+
         # Verify that RabbitMQ responsiveness was checked
         mock_rabbitmq_connection.assert_called_once_with(
             app.state.settings.RABBITMQ.dsn
         )
 
         # Verify that RabbitMQ settings are in the lifespan manager state
         assert app.state.settings.RABBITMQ
+        assert app.state.rpc_server
+        assert app.state.rpc_client
 
     # No explicit shutdown logic for RabbitMQ in this case
-    assert mock_rabbitmq_rpc_client_class.called
+    assert mock_rabbitmq_rpc_client_class.close.called
 
 
 async def test_lifespan_rabbitmq_with_invalid_settings(