✨ Refactor: Ensure async engine disposal in setup_postgres_database with error handling and update test structure

Author: pcrespov

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

@@ -34,7 +34,11 @@ async def setup_postgres_database(_, state: State) -> AsyncIterator[State]:
             pg_settings
         )
 
-        yield {PostgresLifespanStateKeys.POSTGRES_ASYNC_ENGINE: async_engine}
-
-        with log_catch(_logger, reraise=False):
-            await async_engine.dispose()
+        try:
+            yield {
+                PostgresLifespanStateKeys.POSTGRES_ASYNC_ENGINE: async_engine,
+            }
+
+        finally:
+            with log_catch(_logger, reraise=False):
+                await async_engine.dispose()

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

@@ -40,11 +40,11 @@ def app_environment(monkeypatch: pytest.MonkeyPatch) -> EnvVarsDict:
     )
 
 
-async def test_setup_postgres_database_in_an_app(
-    is_pdb_enabled: bool,
+@pytest.fixture
+def app_lifespan(
     app_environment: EnvVarsDict,
     mock_create_async_engine_and_database_ready: MockType,
-):
+) -> LifespanManager:
     assert app_environment
 
     class AppSettings(BaseApplicationSettings):
@@ -60,24 +60,28 @@ async def my_app_settings(app: FastAPI) -> AsyncIterator[State]:
             PostgresLifespanStateKeys.POSTGRES_SETTINGS: app.state.settings.CATALOG_POSTGRES
         }
 
-    async def my_db_setup(app: FastAPI, state: State) -> AsyncIterator[State]:
+    async def my_database_setup(app: FastAPI, state: State) -> AsyncIterator[State]:
         app.state.my_db_engine = state[PostgresLifespanStateKeys.POSTGRES_ASYNC_ENGINE]
 
-        assert (
-            app.state.my_db_engine
-            == mock_create_async_engine_and_database_ready.return_value
-        )
-
         yield {}
 
     # compose lifespans
     app_lifespan = LifespanManager()
     app_lifespan.add(my_app_settings)
 
-    postgres_lifespan.add(my_db_setup)
+    postgres_lifespan.add(my_database_setup)
     app_lifespan.include(postgres_lifespan)
 
-    # define app
+    return app_lifespan
+
+
+async def test_setup_postgres_database_in_an_app(
+    is_pdb_enabled: bool,
+    app_environment: EnvVarsDict,
+    mock_create_async_engine_and_database_ready: MockType,
+    app_lifespan: LifespanManager,
+):
+
     app = FastAPI(lifespan=app_lifespan)
 
     async with ASGILifespanManager(
@@ -103,6 +107,43 @@ async def my_db_setup(app: FastAPI, state: State) -> AsyncIterator[State]:
             ]
         )
 
+        assert (
+            app.state.my_db_engine
+            == mock_create_async_engine_and_database_ready.return_value
+        )
+
     # Verify that the engine was disposed
     async_engine: Any = mock_create_async_engine_and_database_ready.return_value
     async_engine.dispose.assert_called_once()
+
+
+async def test_setup_postgres_database_dispose_engine_on_failure(
+    is_pdb_enabled: bool,
+    app_environment: EnvVarsDict,
+    mock_create_async_engine_and_database_ready: MockType,
+    app_lifespan: LifespanManager,
+):
+    expected_msg = "my_faulty_setup error"
+
+    def raise_error():
+        raise RuntimeError(expected_msg)
+
+    @app_lifespan.add
+    async def my_faulty_setup(app: FastAPI, state: State) -> AsyncIterator[State]:
+        assert PostgresLifespanStateKeys.POSTGRES_ASYNC_ENGINE in state
+        raise_error()
+        yield {}
+
+    app = FastAPI(lifespan=app_lifespan)
+
+    with pytest.raises(RuntimeError, match=expected_msg):
+        async with ASGILifespanManager(
+            app,
+            startup_timeout=None if is_pdb_enabled else 10,
+            shutdown_timeout=None if is_pdb_enabled else 10,
+        ):
+            ...
+
+    # Verify that the engine was disposed even if error happend
+    async_engine: Any = mock_create_async_engine_and_database_ready.return_value
+    async_engine.dispose.assert_called_once()