From 59d2cdd9de748289524376f5fae9fe7c3eaa4186 Mon Sep 17 00:00:00 2001 From: Marat Sarbasov Date: Sun, 16 Jan 2022 04:37:10 +0300 Subject: [PATCH] Add test_dynamic_limit_provider_depending_on_key --- tests/test_fastapi_extension.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/tests/test_fastapi_extension.py b/tests/test_fastapi_extension.py index 60bb72d..5c837a0 100644 --- a/tests/test_fastapi_extension.py +++ b/tests/test_fastapi_extension.py @@ -232,6 +232,33 @@ class TestDecorators(TestSlowapi): r"""parameter `response` must be an instance of starlette.responses.Response""" ) + def test_dynamic_limit_provider_depending_on_key(self): + def custom_key_func(request: Request): + if request.headers.get("TOKEN") == "secret": + return "admin" + return "user" + + def dynamic_limit_provider(key: str): + if key == "admin": + return "10/minute" + return "5/minute" + + app, limiter = self.build_fastapi_app(key_func=custom_key_func) + + @app.get("/t1") + @limiter.limit(dynamic_limit_provider) + async def t1(request: Request, response: Response): + return {"key": "value"} + + client = TestClient(app) + for i in range(0, 10): + response = client.get("/t1") + assert response.status_code == 200 if i < 5 else 429 + + for i in range(0, 20): + response = client.get("/t1", headers={"TOKEN": "secret"}) + assert response.status_code == 200 if i < 10 else 429 + def test_disabled_limiter(self): """ Check that the limiter does nothing if disabled (both sync and async)