From 591f804827cddc8b2d844c2c592adadab9507e8a Mon Sep 17 00:00:00 2001 From: Waket Zheng Date: Mon, 1 Jul 2024 14:13:27 +0800 Subject: [PATCH] feat: add today lock --- README.md | 3 + README.zh.md | 3 + fastapi_cdn_host/__init__.py | 3 +- fastapi_cdn_host/utils.py | 88 +++++++++++++++++++++++---- tests/lock_docs/main.py | 26 +++++++- tests/lock_docs/test_lock.py | 115 +++++++++++++++++++++++++++-------- 6 files changed, 197 insertions(+), 41 deletions(-) diff --git a/README.md b/README.md index 8ecdbca..a00a6b5 100644 --- a/README.md +++ b/README.md @@ -28,6 +28,9 @@ app = FastAPI() fastapi_cdn_host.patch_docs(app) ``` +See more at: +- examples/ +- tests/ ## Detail 1. Let's say that the default docs CDN host https://cdn.jsdelivr.net is too slow in your network, while unpkg.com is much faster. diff --git a/README.zh.md b/README.zh.md index da116f9..c46addb 100644 --- a/README.zh.md +++ b/README.zh.md @@ -27,6 +27,9 @@ app = FastAPI() fastapi_cdn_host.patch_docs(app) ``` +更多示例见： +- examples/ +- tests/ ## 详解 diff --git a/fastapi_cdn_host/__init__.py b/fastapi_cdn_host/__init__.py index c84ea24..0226277 100644 --- a/fastapi_cdn_host/__init__.py +++ b/fastapi_cdn_host/__init__.py @@ -2,7 +2,7 @@ from pathlib import Path from .client import CdnHostEnum, CdnHostItem, monkey_patch_for_docs_ui -from .utils import weekday_lock +from .utils import today_lock, weekday_lock patch_docs = monkey_patch = monkey_patch_for_docs_ui __version__ = importlib.metadata.version(Path(__file__).parent.name) @@ -13,5 +13,6 @@ "monkey_patch", "monkey_patch_for_docs_ui", "patch_docs", + "today_lock", "weekday_lock", ) diff --git a/fastapi_cdn_host/utils.py b/fastapi_cdn_host/utils.py index ff2feb1..2aa1a99 100644 --- a/fastapi_cdn_host/utils.py +++ b/fastapi_cdn_host/utils.py @@ -2,7 +2,7 @@ import sys from contextlib import asynccontextmanager from datetime import datetime -from typing import AsyncGenerator +from typing import AsyncGenerator, Optional from fastapi import HTTPException, Request, status from httpx import ASGITransport, AsyncClient @@ -26,8 +26,68 @@ async def TestClient( yield c +class ParamLock: + param_name: str = "a" + + @staticmethod + def validate(value: str) -> bool: + return True + + @classmethod + def check_param(cls, request: Request, name: str) -> None: + if not (d := request.query_params.get(name)) or not cls.validate(d): + status_code = ( + status.HTTP_418_IM_A_TEAPOT + if sys.version_info >= (3, 9) + else status.HTTP_417_EXPECTATION_FAILED + ) + raise HTTPException(status_code=status_code) + + def __init__(self, name: Optional[str] = None) -> None: + if name is None: + name = self.param_name + self.name = name + + def __call__(self, request: Request) -> None: + self.check_param(request, self.name) + + +class WeekdayLock(ParamLock): + """Check whether docs query param. + + Usage:: + >>> import fastapi_cdn_host + >>> from fastapi import FastAPI + >>> app = FastAPI(openapi_url='/v1/api.json') + >>> fastapi_cdn_host.patch_docs(app, lock=fastapi_cdn_host.utils.WeekdayLock()) + """ + + param_name = "day" + + @staticmethod + def validate(value: str) -> bool: + return (weekday := getattr(calendar, value.upper(), None)) is not None and ( + weekday == datetime.now().weekday() + ) + + +class TodayLock(WeekdayLock): + """Check whether docs query param. + + Usage:: + >>> import fastapi_cdn_host + >>> from fastapi import FastAPI + >>> app = FastAPI(openapi_url='/v1/api.json') + >>> fastapi_cdn_host.patch_docs(app, lock=fastapi_cdn_host.utils.TodayLock()) + """ + + @staticmethod + def validate(value: str) -> bool: + return value == str(datetime.now().date()) + + def weekday_lock(request: Request, name="day") -> None: - """A simple docs lock function. + """Check docs/ query contains `day` param with value of today's weekday, e.g.: Monday. Usage:: >>> import fastapi_cdn_host @@ -35,14 +95,16 @@ def weekday_lock(request: Request, name="day") -> None: >>> app = FastAPI(openapi_url='/v1/api.json') >>> fastapi_cdn_host.patch_docs(app, lock=fastapi_cdn_host.weekday_lock) """ - if ( - not (d := request.query_params.get(name)) - or (weekday := getattr(calendar, d.upper(), None)) is None - or (weekday != datetime.now().weekday()) - ): - status_code = ( - status.HTTP_418_IM_A_TEAPOT - if sys.version_info >= (3, 9) - else status.HTTP_417_EXPECTATION_FAILED - ) - raise HTTPException(status_code=status_code) + WeekdayLock(name)(request) + + +def today_lock(request: Request, name="day") -> None: + """Check docs query param contains `day` and its value is today. + + Usage:: + >>> import fastapi_cdn_host + >>> from fastapi import FastAPI + >>> app = FastAPI(openapi_url='/v1/api.json') + >>> fastapi_cdn_host.patch_docs(app, lock=fastapi_cdn_host.today_lock) + """ + TodayLock(name)(request) diff --git a/tests/lock_docs/main.py b/tests/lock_docs/main.py index 6650610..cccd7ca 100644 --- a/tests/lock_docs/main.py +++ b/tests/lock_docs/main.py @@ -3,12 +3,25 @@ from fastapi import FastAPI, Request import fastapi_cdn_host -from fastapi_cdn_host.utils import weekday_lock +from fastapi_cdn_host.utils import ( + ParamLock, + TodayLock, + WeekdayLock, + today_lock, + weekday_lock, +) app = FastAPI() +app_weekday = FastAPI() +app_weekday_class = FastAPI() app_sync_lock = FastAPI() app_change_lock_param = FastAPI() +app_today = FastAPI() +app_today_class = FastAPI() +app_today_param = FastAPI() +app_param_lock = FastAPI() + def sync_lock(request: Request) -> None: return weekday_lock(request) @@ -19,14 +32,25 @@ async def lock(request: Request) -> None: fastapi_cdn_host.patch_docs(app, lock=lock) +fastapi_cdn_host.patch_docs(app_weekday, lock=weekday_lock) +fastapi_cdn_host.patch_docs(app_weekday_class, lock=WeekdayLock()) fastapi_cdn_host.patch_docs(app_sync_lock, lock=sync_lock) fastapi_cdn_host.patch_docs( app_change_lock_param, lock=partial(fastapi_cdn_host.weekday_lock, name="weekday") ) +fastapi_cdn_host.patch_docs(app_today, lock=fastapi_cdn_host.today_lock) +fastapi_cdn_host.patch_docs(app_today_param, lock=partial(today_lock, name="date")) +fastapi_cdn_host.patch_docs(app_today_class, lock=TodayLock()) +fastapi_cdn_host.patch_docs(app_param_lock, lock=ParamLock()) + @app.get("/") @app_sync_lock.get("/") @app_change_lock_param.get("/") +@app_today.get("/") +@app_today_param.get("/") +@app_today_class.get("/") +@app_param_lock.get("/") async def index(): return "homepage" diff --git a/tests/lock_docs/test_lock.py b/tests/lock_docs/test_lock.py index 41accb8..7756c04 100644 --- a/tests/lock_docs/test_lock.py +++ b/tests/lock_docs/test_lock.py @@ -1,54 +1,46 @@ # mypy: no-disallow-untyped-decorators +import random import sys from datetime import datetime import httpx import pytest -from main import app, app_change_lock_param, app_sync_lock +from main import ( + app, + app_change_lock_param, + app_param_lock, + app_sync_lock, + app_today, + app_today_class, + app_today_param, +) from fastapi_cdn_host.utils import TestClient @pytest.fixture(scope="module") -async def client(): +async def client_app(): async with TestClient(app) as c: yield c @pytest.fixture(scope="module") async def client_sync_lock(): - async with TestClient( - app_sync_lock, - base_url="http://sync.test.com", - ) as c: + async with TestClient(app_sync_lock) as c: yield c @pytest.fixture(scope="module") async def client_change_param_name(): - async with TestClient( - app_change_lock_param, - base_url="http://lock.test.com", - ) as c: + async with TestClient(app_change_lock_param) as c: yield c -class TestLock: - weekdays = [ - "Monday", - "Tuesday", - "Wednesday", - "Thursday", - "Friday", - "Saturday", - "Sunday", - ] - - @pytest.mark.anyio - async def test_lock(self, client: httpx.AsyncClient): - await self.request_locked(client) +class LockTester: + def get_day_value(self) -> str: + return "" - async def request_locked(self, client, param_name="day"): + async def request_locked(self, client: httpx.AsyncClient, param_name="day"): status_code = 418 if sys.version_info >= (3, 9) else 417 response = await client.get("/") assert response.status_code == 200 @@ -59,7 +51,7 @@ async def request_locked(self, client, param_name="day"): response = await client.get("/redoc") assert response.status_code == status_code assert response.json()["detail"] in ("I'm a Teapot", "Expectation Failed") - day = self.weekdays[datetime.now().weekday()] + day = self.get_day_value() response = await client.get(f"/docs?{param_name}={day}") assert response.status_code == 200 response = await client.get(f"/redoc?{param_name}={day}") @@ -68,6 +60,25 @@ async def request_locked(self, client, param_name="day"): assert response.status_code == 200 assert response.text == '"homepage"' + +class TestWeekdayLock(LockTester): + weekdays = [ + "Monday", + "Tuesday", + "Wednesday", + "Thursday", + "Friday", + "Saturday", + "Sunday", + ] + + def get_day_value(self) -> str: + return self.weekdays[datetime.now().weekday()] + + @pytest.mark.anyio + async def test_lock(self, client_app: httpx.AsyncClient): + await self.request_locked(client_app) + @pytest.mark.anyio async def test_sync_lock(self, client_sync_lock: httpx.AsyncClient): await self.request_locked(client_sync_lock) @@ -77,3 +88,55 @@ async def test_change_lock_param_name( self, client_change_param_name: httpx.AsyncClient ): await self.request_locked(client_change_param_name, "weekday") + + +@pytest.fixture(scope="module") +async def client_today(): + async with TestClient(app_today) as c: + yield c + + +@pytest.fixture(scope="module") +async def client_today_class(): + async with TestClient(app_today_class) as c: + yield c + + +@pytest.fixture(scope="module") +async def client_today_param(): + async with TestClient(app_today_param) as c: + yield c + + +@pytest.fixture(scope="module") +async def client_param_lock(): + async with TestClient(app_param_lock) as c: + yield c + + +class TestTodayLock(LockTester): + def get_day_value(self) -> str: + return str(datetime.now().date()) + + @pytest.mark.anyio + async def test_today_lock(self, client_today: httpx.AsyncClient): + await self.request_locked(client_today) + + @pytest.mark.anyio + async def test_today_class(self, client_today_class: httpx.AsyncClient): + await self.request_locked(client_today_class) + + @pytest.mark.anyio + async def test_today_param(self, client_today_param: httpx.AsyncClient): + await self.request_locked(client_today_param, "date") + + +class TestParamLock(LockTester): + def get_day_value(self) -> str: + return str(random.randint(1, 100)) + + @pytest.mark.anyio + async def test_param_lock(self, client_param_lock: httpx.AsyncClient): + await self.request_locked(client_param_lock, "a") + await self.request_locked(client_param_lock, "a") + await self.request_locked(client_param_lock, "a")