Add wrapper for app lifespan

fastapi_pagination/api.py

@@ -14,7 +14,7 @@
 
 import inspect
 import warnings
-from contextlib import ExitStack, contextmanager, suppress
+from contextlib import ExitStack, asynccontextmanager, contextmanager, suppress
 from contextvars import ContextVar
 from typing import (
     Any,
@@ -365,8 +365,15 @@ def _add_pagination(parent: ParentT) -> None:
 def add_pagination(parent: ParentT) -> ParentT:
     _add_pagination(parent)
 
-    @parent.on_event("startup")
-    def on_startup() -> None:
+    router = parent.router if isinstance(parent, FastAPI) else parent
+    _original_lifespan_context = router.lifespan_context
+
+    @asynccontextmanager
+    async def lifespan(app: Any) -> AsyncIterator[Any]:
         _add_pagination(parent)
 
+        async with _original_lifespan_context(app) as maybe_state:
+            yield maybe_state
+
+    router.lifespan_context = lifespan
     return parent

tests/test_api.py

@@ -1,9 +1,11 @@
-from typing import Any, Generic, Sequence, TypeVar
+from contextlib import asynccontextmanager
+from typing import Any, AsyncIterator, Generic, Sequence, TypeVar
 
 from fastapi import Depends, FastAPI, Request, Response, status
 from fastapi.routing import APIRouter
 from fastapi.testclient import TestClient
 from pydantic import Field
+from pytest import fixture
 
 try:
     from pydantic.generics import GenericModel
@@ -252,3 +254,32 @@ async def route():
 
     rsp = client.get("/", params={"page": -1})
     assert rsp.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY
+
+
+class TestLifespan:
+    @fixture(scope="class")
+    def app(self):
+        @asynccontextmanager
+        async def lifespan(_: Any) -> AsyncIterator[None]:
+            app.state.called = True
+            yield
+
+        app = FastAPI(lifespan=lifespan)
+
+        @app.get(
+            "/",
+            response_model=Page[int],
+        )
+        async def route():
+            return paginate([])
+
+        add_pagination(app)
+        app.state.called = False
+
+        return app
+
+    async def test_lifespan_wrap(self, app, client):
+        rsp = await client.get("/")
+        assert rsp.status_code == status.HTTP_200_OK
+
+        assert app.state.called, "original lifespan not called"