python/slowapi
1
0
Fork 0
mirror of https://github.com/laurentS/slowapi.git synced 2026-03-13 09:10:20 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix formatting

Browse Source
This commit is contained in:
Guillaume Gardey
2020-12-23 10:04:52 +00:00
parent 755bf44e12
commit b6fa460254
2 changed files with 15 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
slowapi/extension.py
View File

@@ -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

15
tests/test_fastapi_extension.py
View File

@@ -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)