python-learning/rate_limiters/tests/test_middleware_starlette.py

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