diff --git a/slowapi/extension.py b/slowapi/extension.py index 0573c72..c0db72a 100644 --- a/slowapi/extension.py +++ b/slowapi/extension.py @@ -629,7 +629,9 @@ class Limiter: response = await func(*args, **kwargs) # type: ignore if not isinstance(response, Response): # get the response object from the decorated endpoint function - self._inject_headers(kwargs.get("response"), request.state.view_rate_limit) + self._inject_headers( + kwargs.get("response"), request.state.view_rate_limit + ) else: self._inject_headers(response, request.state.view_rate_limit) return response @@ -655,7 +657,9 @@ class Limiter: response = func(*args, **kwargs) if not isinstance(response, Response): # get the response object from the decorated endpoint function - self._inject_headers(kwargs.get("response"), request.state.view_rate_limit) + self._inject_headers( + kwargs.get("response"), request.state.view_rate_limit + ) else: self._inject_headers(response, request.state.view_rate_limit) return response diff --git a/tests/test_fastapi_extension.py b/tests/test_fastapi_extension.py index 7580450..6f2ef94 100644 --- a/tests/test_fastapi_extension.py +++ b/tests/test_fastapi_extension.py @@ -22,7 +22,6 @@ class TestDecorators(TestSlowapi): response = client.get("/t1") assert response.status_code == 200 if i < 5 else 429 - def test_single_decorator_with_headers(self): app, limiter = self.build_fastapi_app(key_func=get_ipaddr, headers_enabled=True) @@ -35,8 +34,10 @@ class TestDecorators(TestSlowapi): for i in range(0, 10): response = client.get("/t1") assert response.status_code == 200 if i < 5 else 429 - assert response.headers.get('X-RateLimit-Limit') is not None if i < 5 else True - assert response.headers.get('Retry-After') is not None if i < 5 else True + assert ( + response.headers.get("X-RateLimit-Limit") is not None if i < 5 else True + ) + assert response.headers.get("Retry-After") is not None if i < 5 else True def test_single_decorator_not_response(self): app, limiter = self.build_fastapi_app(key_func=get_ipaddr) @@ -63,8 +64,10 @@ class TestDecorators(TestSlowapi): for i in range(0, 10): response = client.get("/t1") assert response.status_code == 200 if i < 5 else 429 - assert response.headers.get('X-RateLimit-Limit') is not None if i < 5 else True - assert response.headers.get('Retry-After') is not None if i < 5 else True + assert ( + response.headers.get("X-RateLimit-Limit") is not None if i < 5 else True + ) + assert response.headers.get("Retry-After") is not None if i < 5 else True def test_multiple_decorators(self): app, limiter = self.build_fastapi_app(key_func=get_ipaddr) @@ -100,7 +103,7 @@ class TestDecorators(TestSlowapi): ) # effectively becomes a limit for all users @limiter.limit("50/minute") # per ip as per default key_func async def t1(request: Request, response: Response): - return {"key":"value"} + return {"key": "value"} with hiro.Timeline().freeze() as timeline: cli = TestClient(app)