50 lines
1.1 KiB
Python
50 lines
1.1 KiB
Python
from time import sleep
|
|
|
|
from starlette.applications import Starlette
|
|
from starlette.testclient import TestClient
|
|
|
|
from middleware_starlette import app, RateLimiterMiddleware, routes
|
|
|
|
def test_app_runs():
|
|
client = TestClient(app)
|
|
|
|
response = client.get("/")
|
|
|
|
assert response.status_code == 200
|
|
assert "Hello there" in response.text
|
|
|
|
|
|
def test_below_rate_limit():
|
|
client = TestClient(app)
|
|
|
|
# Limit is 10/s
|
|
for i in range(20):
|
|
response = client.get("/")
|
|
|
|
assert response.status_code == 200
|
|
assert "Hello there" in response.text
|
|
sleep(0.1)
|
|
|
|
|
|
def test_above_rate_limit():
|
|
client = TestClient(app)
|
|
|
|
successful = 0
|
|
rate_limited = 0
|
|
|
|
# Limit is 10/s
|
|
for i in range(30):
|
|
response = client.get("/")
|
|
|
|
if response.status_code != 200:
|
|
rate_limited += 1
|
|
assert response.status_code == 429
|
|
assert "Rate limit exceeded" in response.text
|
|
assert response.headers.get("retry-after") is not None
|
|
else:
|
|
successful += 1
|
|
|
|
assert successful != 0
|
|
assert rate_limited != 0
|
|
|