added combine_lfiespan_context_managers

Author: Andrei Neagu

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

@@ -0,0 +1,24 @@
+from collections.abc import AsyncGenerator, Callable
+from contextlib import AsyncExitStack, asynccontextmanager
+from typing import AsyncContextManager, TypeAlias
+
+from fastapi import FastAPI
+
+LifespanContextManager: TypeAlias = Callable[[FastAPI], AsyncContextManager[None]]
+
+
+def combine_lfiespan_context_managers(
+    *context_managers: LifespanContextManager,
+) -> LifespanContextManager:
+    """the first entry has its `setup` called first and its `teardown` called last
+    With `setup` and `teardown` referring to the code before and after the `yield`
+    """
+
+    @asynccontextmanager
+    async def _(app: FastAPI) -> AsyncGenerator[None, None]:
+        async with AsyncExitStack() as stack:
+            for context_manager in context_managers:
+                await stack.enter_async_context(context_manager(app))
+            yield
+
+    return _

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

@@ -0,0 +1,42 @@
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+
+import pytest
+from asgi_lifespan import LifespanManager
+from fastapi import FastAPI
+from servicelib.fastapi.lifespan_utils import combine_lfiespan_context_managers
+
+
+async def test_multiple_lifespan_managers(capsys: pytest.CaptureFixture):
+    @asynccontextmanager
+    async def database_lifespan(_: FastAPI) -> AsyncIterator[None]:
+        print("setup DB")
+        yield
+        print("shutdown  DB")
+
+    @asynccontextmanager
+    async def cache_lifespan(_: FastAPI) -> AsyncIterator[None]:
+        print("setup CACHE")
+        yield
+        print("shutdown CACHE")
+
+    app = FastAPI(
+        lifespan=combine_lfiespan_context_managers(database_lifespan, cache_lifespan)
+    )
+
+    capsys.readouterr()
+
+    async with LifespanManager(app):
+        messages = capsys.readouterr().out
+
+        assert "setup DB" in messages
+        assert "setup CACHE" in messages
+        assert "shutdown  DB" not in messages
+        assert "shutdown CACHE" not in messages
+
+    messages = capsys.readouterr().out
+
+    assert "setup DB" not in messages
+    assert "setup CACHE" not in messages
+    assert "shutdown  DB" in messages
+    assert "shutdown CACHE" in messages