add override_defaults parameter and fix merge test

This commit is contained in:
Reuben Thomas-Davis
2020-10-01 19:05:13 +01:00
parent 272ae50d05
commit b8fa36ea83
3 changed files with 40 additions and 18 deletions

View File

@@ -70,8 +70,8 @@ class TestDecorators(TestSlowapi):
def test_multiple_decorators(self):
app, limiter = self.build_starlette_app(key_func=get_ipaddr)
@limiter.limit("100 per minute", lambda: "test")
@limiter.limit("50/minute") # per ip as per default key_func
@limiter.limit("10 per minute", lambda: "test")
@limiter.limit("5/minute") # per ip as per default key_func
async def t1(request: Request):
return PlainTextResponse("test")
@@ -79,16 +79,16 @@ class TestDecorators(TestSlowapi):
with hiro.Timeline().freeze() as timeline:
cli = TestClient(app)
for i in range(0, 100):
for i in range(0, 10):
response = cli.get("/t1", headers={"X_FORWARDED_FOR": "127.0.0.2"})
assert response.status_code == 200 if i < 50 else 429
for i in range(50):
assert response.status_code == 200 if i < 5 else 429
for i in range(5):
assert cli.get("/t1").status_code == 200
assert cli.get("/t1").status_code == 429
assert (
cli.get("/t1", headers={"X_FORWARDED_FOR": "127.0.0.3"}).status_code
== 429
cli.get("/t1", headers={"X_FORWARDED_FOR": "127.0.0.3"}).status_code
== 429
)
def test_headers_no_breach(self):
@@ -196,10 +196,13 @@ class TestDecorators(TestSlowapi):
resp2 = cli.get("/t2", headers={"X_FORWARDED_FOR": "127.0.0.10"})
assert resp2.status_code == 200
# todo: more tests - see https://github.com/alisaifee/flask-limiter/blob/55df08f14143a7e918fc033067a494248ab6b0c5/tests/test_decorators.py#L187
def test_default_and_decorator_limit_merging(self):
app, limiter = self.build_starlette_app(key_func=get_ipaddr, default_limits=["50/minute"])
# test pool has 100 reqs left
app, limiter = self.build_starlette_app(key_func=lambda: "test", default_limits=["10/minute"])
@limiter.limit("100 per minute", key_func=lambda: "lest")
# ip pool has 50 reqs for 127.0.0.14
@limiter.limit("5 per minute", key_func=get_ipaddr, override_defaults=False)
async def t1(request: Request):
return PlainTextResponse("test")
@@ -207,13 +210,15 @@ class TestDecorators(TestSlowapi):
with hiro.Timeline().freeze() as timeline:
cli = TestClient(app)
for i in range(0, 100):
response = cli.get("/t1", headers={"X_FORWARDED_FOR": "127.0.0.14"})
assert response.status_code == 200 if i < 50 else 429
for i in range(50):
for i in range(0, 10):
response = cli.get("/t1", headers={"X_FORWARDED_FOR": "127.0.0.2"})
assert response.status_code == 200 if i < 5 else 429
# now ip pool for 127.0.0.14 has 0 reqs left
for i in range(5):
assert cli.get("/t1").status_code == 200
# now test pool has 0 reqs left
assert cli.get("/t1").status_code == 429
assert (
cli.get("/t1", headers={"X_FORWARDED_FOR": "127.0.0.15"}).status_code
cli.get("/t1", headers={"X_FORWARDED_FOR": "127.0.0.3"}).status_code
== 429)