Add methods for set cache_age in openapi schema

fast_cache_middleware/_helpers.py

@@ -0,0 +1,28 @@
+from fastapi import FastAPI, routing
+
+from .depends import CacheConfig
+
+
+def set_cache_age_in_openapi_schema(app: FastAPI) -> None:
+    openapi_schema = app.openapi()
+
+    for route in app.routes:
+        if isinstance(route, routing.APIRoute):
+            path = route.path
+            methods = route.methods
+
+            for dependency in route.dependencies:
+                dep = dependency.dependency
+                if isinstance(dep, CacheConfig):
+                    max_age = dep.max_age
+
+                    for method in methods:
+                        method = method.lower()
+                        try:
+                            operation = openapi_schema["paths"][path][method]
+                            operation.setdefault("x-cache-age", max_age)
+                        except KeyError:
+                            continue
+
+    app.openapi_schema = openapi_schema
+    return None

fast_cache_middleware/middleware.py

@@ -8,6 +8,7 @@
 from starlette.routing import Match, Mount
 from starlette.types import ASGIApp, Receive, Scope, Send
 
+from ._helpers import set_cache_age_in_openapi_schema
 from .controller import Controller
 from .depends import BaseCacheConfigDepends, CacheConfig, CacheDropConfig
 from .schemas import RouteInfo
@@ -217,6 +218,7 @@ def __init__(
 
         self.storage = storage or InMemoryStorage()
         self.controller = controller or Controller()
+        self._openapi_initialized = False
 
         self._routes_info: list[RouteInfo] = []
 
@@ -229,12 +231,17 @@ def __init__(
 
     async def on_lifespan(self, scope: Scope, _: Receive, __: Send) -> bool | None:
         app_routes = get_app_routes(scope["app"])
+        set_cache_age_in_openapi_schema(scope["app"])
         self._routes_info = self._extract_routes_info(app_routes)
         return None
 
     async def on_http(self, scope: Scope, receive: Receive, send: Send) -> bool | None:
         request = Request(scope, receive)
 
+        if not self._openapi_initialized:
+            set_cache_age_in_openapi_schema(scope["app"])
+            self._openapi_initialized = True
+
         # Find matching route
         route_info = self._find_matching_route(request, self._routes_info)
         if not route_info:

tests/conftest.py

@@ -0,0 +1,141 @@
+import time
+import typing as tp
+from http import HTTPMethod
+from types import MethodType
+
+import pytest
+from fastapi import Depends, FastAPI, HTTPException
+from fastapi.openapi.utils import get_openapi
+from starlette.requests import Request
+from starlette.testclient import TestClient
+
+from fast_cache_middleware import CacheConfig, CacheDropConfig, FastCacheMiddleware
+
+
+async def get_storage_depends(request: Request) -> tp.Any:
+    return request.app.state.storage
+
+
+async def create_user(user_id: int, storage: dict = Depends(get_storage_depends)):
+    user_name = str(time.time)
+    storage[user_id] = user_name
+    return {"user_id": user_id, "name": user_name, "timestamp": time.time()}
+
+
+async def get_user(user_id: int, storage: dict = Depends(get_storage_depends)):
+    try:
+        user_name = storage[user_id]
+    except KeyError:
+        raise HTTPException(status_code=404, detail="User not found")
+    return {"user_id": user_id, "name": user_name, "timestamp": time.time()}
+
+
+async def delete_user(user_id: int, storage: dict = Depends(get_storage_depends)):
+    user_name = storage.pop(user_id)
+    return {"user_id": user_id, "name": user_name, "timestamp": time.time()}
+
+
+async def get_first_user(storage: dict = Depends(get_storage_depends)):
+    return await get_user(1, storage=storage)
+
+
+async def get_second_user(storage: dict = Depends(get_storage_depends)):
+    return await get_user(2, storage=storage)
+
+
+async def hidden_route():
+    return {"hidden": True}
+
+
+@pytest.fixture
+def app() -> FastAPI:
+    _storage = {
+        1: "first",
+        2: "second",
+    }
+
+    app = FastAPI()
+    app.add_middleware(FastCacheMiddleware)
+    app.state.storage = _storage
+
+    app.router.add_api_route(
+        "/users/first",
+        get_first_user,
+        dependencies=[CacheDropConfig(["/users/second"])],
+        methods={HTTPMethod.GET.value},
+    )
+    app.router.add_api_route(
+        "/users/second",
+        get_second_user,
+        dependencies=[CacheConfig(max_age=5)],
+        methods={HTTPMethod.GET.value},
+    )
+    app.router.add_api_route(
+        "/users/{user_id}",
+        get_user,
+        dependencies=[CacheConfig(max_age=10)],
+        methods={HTTPMethod.GET.value},
+    )
+    app.router.add_api_route(
+        "/users/{user_id}",
+        create_user,
+        dependencies=[CacheDropConfig(paths=["/users/"])],
+        methods={HTTPMethod.POST.value},
+    )
+    app.router.add_api_route(
+        "/users/{user_id}",
+        delete_user,
+        dependencies=[CacheDropConfig(paths=["/users/"])],
+        methods={HTTPMethod.DELETE.value},
+    )
+    app.router.add_api_route(
+        "/no-docs",
+        hidden_route,
+        include_in_schema=False,
+        dependencies=[CacheConfig(max_age=42)],
+        methods={HTTPMethod.GET.value},
+    )
+
+    def custom_openapi() -> dict[str, tp.Any]:
+        if app.openapi_schema:
+            return app.openapi_schema
+        openapi_schema = get_openapi(
+            title="Custom title",
+            version="2.5.0",
+            summary="This is a very custom OpenAPI schema",
+            description="Here's a longer description of the custom **OpenAPI** schema",
+            routes=app.routes,
+        )
+        openapi_schema["info"]["x-logo"] = {
+            "url": "https://fastapi.tiangolo.com/img/logo-margin/logo-teal.png"
+        }
+        app.openapi_schema = openapi_schema
+        return app.openapi_schema
+
+    app.openapi = custom_openapi  # type: ignore
+
+    second_app = FastAPI()
+    second_app.add_middleware(FastCacheMiddleware)
+    second_app.state.storage = _storage
+
+    second_app.router.add_api_route(
+        "/users/first",
+        get_first_user,
+        dependencies=[CacheDropConfig(["/subapp/users/second"])],
+        methods={HTTPMethod.GET.value},
+    )
+    second_app.router.add_api_route(
+        "/users/second",
+        get_second_user,
+        dependencies=[CacheConfig(max_age=5)],
+        methods={HTTPMethod.GET.value},
+    )
+    app.mount("/subapp", app=second_app)
+
+    return app
+
+
+@pytest.fixture
+def client(app: FastAPI) -> TestClient:
+    """Создает тестовый клиент."""
+    return TestClient(app)

tests/test_controller.py

@@ -3,7 +3,7 @@
 import asyncio
 import time
 import typing as tp
-from datetime import datetime, timedelta
+from datetime import UTC, datetime
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
@@ -163,7 +163,7 @@ async def test_get_cached_response_success(
         """Тестирует успешное получение кешированного ответа."""
         cached_response = Response(content="cached", status_code=200)
         metadata = {
-            "cached_at": datetime.utcnow().isoformat(),
+            "cached_at": datetime.now(UTC).isoformat(),
             "ttl": 300,
             "etag": "test-etag",
         }

tests/test_middleware.py

@@ -1,115 +1,7 @@
 """Тесты для оптимизированного FastCacheMiddleware."""
 
-import time
-import typing as tp
-from http import HTTPMethod
-
-import pytest
-from fastapi import Depends, FastAPI, HTTPException, Request
 from fastapi.testclient import TestClient
 
-from fast_cache_middleware import CacheConfig, CacheDropConfig, FastCacheMiddleware
-
-
-async def get_storage_depends(request: Request) -> tp.Any:
-    return request.app.state.storage
-
-
-async def create_user(user_id: int, storage: dict = Depends(get_storage_depends)):
-    user_name = str(time.time)
-    storage[user_id] = user_name
-    return {"user_id": user_id, "name": user_name, "timestamp": time.time()}
-
-
-async def get_user(user_id: int, storage: dict = Depends(get_storage_depends)):
-    try:
-        user_name = storage[user_id]
-    except KeyError:
-        raise HTTPException(status_code=404, detail="User not found")
-    return {"user_id": user_id, "name": user_name, "timestamp": time.time()}
-
-
-async def delete_user(user_id: int, storage: dict = Depends(get_storage_depends)):
-    user_name = storage.pop(user_id)
-    return {"user_id": user_id, "name": user_name, "timestamp": time.time()}
-
-
-async def get_first_user(storage: dict = Depends(get_storage_depends)):
-    return await get_user(1, storage=storage)
-
-
-async def get_second_user(storage: dict = Depends(get_storage_depends)):
-    return await get_user(2, storage=storage)
-
-
-@pytest.fixture
-def app() -> FastAPI:
-    _storage = {
-        1: "first",
-        2: "second",
-    }
-
-    app = FastAPI()
-    app.add_middleware(FastCacheMiddleware)
-    app.state.storage = _storage
-
-    app.router.add_api_route(
-        "/users/first",
-        get_first_user,
-        dependencies=[CacheDropConfig(["/users/second"])],
-        methods={HTTPMethod.GET.value},
-    )
-    app.router.add_api_route(
-        "/users/second",
-        get_second_user,
-        dependencies=[CacheConfig(max_age=5)],
-        methods={HTTPMethod.GET.value},
-    )
-    app.router.add_api_route(
-        "/users/{user_id}",
-        get_user,
-        dependencies=[CacheConfig(max_age=10)],
-        methods={HTTPMethod.GET.value},
-    )
-    app.router.add_api_route(
-        "/users/{user_id}",
-        create_user,
-        dependencies=[CacheDropConfig(paths=["/users/"])],
-        methods={HTTPMethod.POST.value},
-    )
-    app.router.add_api_route(
-        "/users/{user_id}",
-        delete_user,
-        dependencies=[CacheDropConfig(paths=["/users/"])],
-        methods={HTTPMethod.DELETE.value},
-    )
-
-    second_app = FastAPI()
-    second_app.add_middleware(FastCacheMiddleware)
-    second_app.state.storage = _storage
-
-    second_app.router.add_api_route(
-        "/users/first",
-        get_first_user,
-        dependencies=[CacheDropConfig(["/subapp/users/second"])],
-        methods={HTTPMethod.GET.value},
-    )
-    second_app.router.add_api_route(
-        "/users/second",
-        get_second_user,
-        dependencies=[CacheConfig(max_age=5)],
-        methods={HTTPMethod.GET.value},
-    )
-    app.mount("/subapp", app=second_app)
-
-    return app
-
-
-@pytest.fixture
-def client(app: FastAPI) -> TestClient:
-    """Создает тестовый клиент."""
-    return TestClient(app)
-
 
 def test_caching_works(client: TestClient) -> None:
     """Тестирует кеширование"""

tests/test_openapi_schema.py

@@ -0,0 +1,42 @@
+from fastapi import FastAPI
+from starlette.testclient import TestClient
+
+
+def test_set_cache_age_to_openapi_schema(app: FastAPI, client: TestClient) -> None:
+    path = "/users/second"
+    method = "get"
+
+    client.get(path)
+    schema = app.openapi()
+
+    assert path in schema["paths"]
+    assert method in schema["paths"][path]
+
+    operation = schema["paths"][path][method]
+
+    assert "x-cache-age" in operation
+    assert operation["x-cache-age"] == 5
+
+
+def test_x_logo_field_exists_after_set_cache_age(
+    app: FastAPI, client: TestClient
+) -> None:
+    path = "/users/second"
+    method = "get"
+
+    client.get(path)
+    schema = app.openapi()
+    operation = schema["paths"][path][method]
+
+    assert "x-logo" in schema["info"]
+    assert "x-cache-age" in operation
+
+
+def test_openapi_patch_keyerror_handled_gracefully(
+    app: FastAPI, client: TestClient
+) -> None:
+    path = "/no-docs"
+
+    client.get(path)
+    schema = app.openapi()
+    assert path not in schema["paths"]