ratelimit: tests (compteurs Lua, middleware, usage) + sondes /v1/_test et fixtures partagées
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
5 changed files +452 −0
modified
tests/conftest.py
+34 −0
@@ -39,6 +39,8 @@ def app(lake): | ||
| 39 | 39 | for m in [m for m in list(sys.modules) if m.split(".")[0] in ("core", "main", "futures", "accounts", "ratelimit", "fundamentals", "bulk", "stream", "openapi")]: |
| 40 | 40 | del sys.modules[m] |
| 41 | 41 | main = importlib.import_module("main") |
| 42 | + from tests.fixtures.v2_test_routes import install as _install_probes | |
| 43 | + _install_probes(main.app) # /v1/_test/* probes for the rate-limit contract (test app only) | |
| 42 | 44 | return main.app |
| 43 | 45 | |
| 44 | 46 | |
@@ -47,3 +49,35 @@ def client(app): | ||
| 47 | 49 | from fastapi.testclient import TestClient |
| 48 | 50 | with TestClient(app) as c: |
| 49 | 51 | yield c |
| 52 | + | |
| 53 | + | |
| 54 | +@pytest.fixture(autouse=True) | |
| 55 | +def _fresh_counters(): | |
| 56 | + """Every test starts with empty rate-limit / usage counters and an empty key cache (fakeredis only).""" | |
| 57 | + yield | |
| 58 | + if "ratelimit.redis_limiter" in sys.modules: | |
| 59 | + sys.modules["ratelimit.redis_limiter"].reset_for_tests() | |
| 60 | + if "ratelimit.middleware" in sys.modules: | |
| 61 | + sys.modules["ratelimit.middleware"].invalidate_key_cache() | |
| 62 | + | |
| 63 | + | |
| 64 | +@pytest.fixture | |
| 65 | +def make_user(app): | |
| 66 | + """Create an active, verified user with one active API key. Returns (user_id, raw_key, email, password).""" | |
| 67 | + from accounts import service | |
| 68 | + from core.db import session | |
| 69 | + counter = {"n": 0} | |
| 70 | + | |
| 71 | + def _make(email: str | None = None, *, tier: str = "free", role: str = "user", password: str = "correct-horse-battery", | |
| 72 | + with_key: bool = True): | |
| 73 | + counter["n"] += 1 | |
| 74 | + email = email or f"user{counter['n']}-{os.getpid()}-{id(counter)}@example.com" | |
| 75 | + with session() as s: | |
| 76 | + u = service.create_user(s, email, "Test User", password=password, tier=tier, role=role, status="active", actor="test") | |
| 77 | + u.email_verified_at = service.now() | |
| 78 | + raw = None | |
| 79 | + if with_key: | |
| 80 | + raw, _ = service.create_key(s, u, "default", actor="test", notify=False) | |
| 81 | + uid = u.id | |
| 82 | + return uid, raw, email, password | |
| 83 | + return _make | |
added
tests/fixtures/v2_test_routes.py
+43 −0
@@ -0,0 +1,43 @@ | ||
| 1 | +"""Probe endpoints mounted ONLY in the test app (under /v1/_test/) to exercise the rate-limit contract: | |
| 2 | +row accounting (json/csv/parquet), quota_exempt, request_cost, requires_key, key-required tags, clamp_limit. | |
| 3 | +""" | |
| 4 | +from __future__ import annotations | |
| 5 | + | |
| 6 | +import numpy as np | |
| 7 | +import pandas as pd | |
| 8 | +from fastapi import FastAPI, Query, Request | |
| 9 | + | |
| 10 | +from core.responses import clamp_limit, frame_response, json_response, parse_format | |
| 11 | + | |
| 12 | + | |
| 13 | +def install(app: FastAPI) -> None: | |
| 14 | + def frame(request: Request, n: int = Query(10), format: str = Query("json"), limit: int | None = Query(None)): | |
| 15 | + lim = clamp_limit(limit, 10_000, 1_000_000, request) if limit is not None else n | |
| 16 | + rows = min(n, lim) | |
| 17 | + df = pd.DataFrame({"i": np.arange(rows), "x": np.linspace(0, 1, rows) if rows else []}) | |
| 18 | + return frame_response(df, parse_format(format), request=request) | |
| 19 | + | |
| 20 | + def exempt(request: Request, n: int = Query(1000)): | |
| 21 | + request.state.quota_exempt = True | |
| 22 | + return json_response([{"i": i} for i in range(n)]) | |
| 23 | + | |
| 24 | + def expensive(request: Request): | |
| 25 | + request.state.request_cost = 2 | |
| 26 | + return json_response({"ok": True}) | |
| 27 | + | |
| 28 | + def needkey(request: Request): | |
| 29 | + request.state.requires_key = True | |
| 30 | + return json_response({"secret": True}) | |
| 31 | + | |
| 32 | + def tagged(): | |
| 33 | + return json_response({"stream": True}) | |
| 34 | + | |
| 35 | + def boom(): | |
| 36 | + raise RuntimeError("kaboom") | |
| 37 | + | |
| 38 | + app.router.add_api_route("/v1/_test/frame", frame, methods=["GET"], include_in_schema=False) | |
| 39 | + app.router.add_api_route("/v1/_test/exempt", exempt, methods=["GET"], include_in_schema=False) | |
| 40 | + app.router.add_api_route("/v1/_test/expensive", expensive, methods=["GET"], include_in_schema=False) | |
| 41 | + app.router.add_api_route("/v1/_test/needkey", needkey, methods=["GET"], include_in_schema=False) | |
| 42 | + app.router.add_api_route("/v1/_test/tagged", tagged, methods=["GET"], tags=["stream"], include_in_schema=False) | |
| 43 | + app.router.add_api_route("/v1/_test/boom", boom, methods=["GET"], include_in_schema=False) | |
added
tests/test_ratelimit_lua.py
+112 −0
@@ -0,0 +1,112 @@ | ||
| 1 | +"""Lua sliding-window counters (fakeredis + lupa): remaining, reset, window rollover, force, peek, fail-open.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +import pytest | |
| 5 | + | |
| 6 | + | |
| 7 | +@pytest.fixture | |
| 8 | +def rl(app): | |
| 9 | + from ratelimit import redis_limiter | |
| 10 | + redis_limiter.reset_for_tests() | |
| 11 | + return redis_limiter | |
| 12 | + | |
| 13 | + | |
| 14 | +@pytest.fixture | |
| 15 | +def tiers(app): | |
| 16 | + from ratelimit import tiers | |
| 17 | + return tiers | |
| 18 | + | |
| 19 | + | |
| 20 | +T0 = 1_800_000_000_000 # ms | |
| 21 | + | |
| 22 | + | |
| 23 | +def test_requests_counter_and_reset(rl, tiers): | |
| 24 | + keyless = tiers.TIERS["keyless"] | |
| 25 | + d = rl.apply_tier("ip:test", keyless, req_cost=1, at_ms=T0) | |
| 26 | + assert d.allowed and d.applied | |
| 27 | + assert d.remaining_requests == 29 and d.limit_requests == 30 | |
| 28 | + assert d.remaining_rows == 100_000 | |
| 29 | + assert d.reset_requests == T0 // 1000 + 3600 # oldest bucket + window | |
| 30 | + for _ in range(29): | |
| 31 | + d = rl.apply_tier("ip:test", keyless, req_cost=1, at_ms=T0 + 5_000) | |
| 32 | + assert d.allowed and d.remaining_requests == 0 | |
| 33 | + d = rl.apply_tier("ip:test", keyless, req_cost=1, at_ms=T0 + 6_000) | |
| 34 | + assert not d.allowed and not d.allowed_requests and d.allowed_rows and not d.applied | |
| 35 | + assert d.remaining_requests == 0 | |
| 36 | + assert d.reset_requests == T0 // 1000 + 3600 | |
| 37 | + | |
| 38 | + | |
| 39 | +def test_window_rollover_frees_oldest_bucket(rl, tiers): | |
| 40 | + free = tiers.TIERS["free"] # 60 s window, 120 requests | |
| 41 | + for i in range(120): | |
| 42 | + d = rl.apply_tier("key:1", free, req_cost=1, at_ms=T0 + i * 100) # all within 12 s | |
| 43 | + assert d.remaining_requests == 0 | |
| 44 | + assert not rl.apply_tier("key:1", free, req_cost=1, at_ms=T0 + 30_000).allowed | |
| 45 | + # 60 s after the first bucket the first second (10 requests) has rolled out of the window | |
| 46 | + d = rl.apply_tier("key:1", free, req_cost=0, at_ms=T0 + 60_000) | |
| 47 | + assert d.remaining_requests == 10 | |
| 48 | + d = rl.apply_tier("key:1", free, req_cost=1, at_ms=T0 + 60_000) | |
| 49 | + assert d.allowed and d.remaining_requests == 9 | |
| 50 | + # far in the future: everything expired, hash pruned | |
| 51 | + d = rl.apply_tier("key:1", free, req_cost=0, at_ms=T0 + 200_000) | |
| 52 | + assert d.remaining_requests == 120 | |
| 53 | + assert rl.client().hlen("rl:key:1:req") == 0 | |
| 54 | + | |
| 55 | + | |
| 56 | +def test_rows_counter_precheck_and_force(rl, tiers): | |
| 57 | + keyless = tiers.TIERS["keyless"] | |
| 58 | + # charge rows after the fact, even when it overshoots | |
| 59 | + d = rl.apply_tier("ip:r", keyless, req_cost=0, rows_cost=99_990, force=True, at_ms=T0) | |
| 60 | + assert d.applied and d.remaining_rows == 10 | |
| 61 | + d = rl.apply_tier("ip:r", keyless, req_cost=0, rows_cost=50, force=True, at_ms=T0 + 1000) | |
| 62 | + assert d.applied and not d.allowed_rows and d.remaining_rows == 0 | |
| 63 | + # now the pre-check refuses on rows while requests are still available | |
| 64 | + d = rl.apply_tier("ip:r", keyless, req_cost=1, at_ms=T0 + 2000) | |
| 65 | + assert d.allowed_requests and not d.allowed_rows and not d.applied | |
| 66 | + assert d.reset_rows == T0 // 1000 + 3600 | |
| 67 | + | |
| 68 | + | |
| 69 | +def test_peek_does_not_write(rl, tiers): | |
| 70 | + free = tiers.TIERS["free"] | |
| 71 | + d = rl.peek_tier("key:9", free, at_ms=T0) | |
| 72 | + assert d.remaining_requests == 120 and d.remaining_rows == 1_000_000 | |
| 73 | + assert rl.client().exists("rl:key:9:req") == 0 | |
| 74 | + | |
| 75 | + | |
| 76 | +def test_two_principals_are_independent(rl, tiers): | |
| 77 | + keyless = tiers.TIERS["keyless"] | |
| 78 | + rl.apply_tier("ip:a", keyless, at_ms=T0) | |
| 79 | + assert rl.peek_tier("ip:b", keyless, at_ms=T0).remaining_requests == 30 | |
| 80 | + assert rl.peek_tier("ip:a", keyless, at_ms=T0).remaining_requests == 29 | |
| 81 | + | |
| 82 | + | |
| 83 | +def test_fail_open_when_redis_raises(rl, tiers, monkeypatch, caplog): | |
| 84 | + import redis as redis_lib | |
| 85 | + | |
| 86 | + class Broken: | |
| 87 | + def evalsha(self, *a, **k): | |
| 88 | + raise redis_lib.exceptions.ConnectionError("down") | |
| 89 | + | |
| 90 | + def script_load(self, *a): | |
| 91 | + raise redis_lib.exceptions.ConnectionError("down") | |
| 92 | + | |
| 93 | + monkeypatch.setattr(rl, "_client", Broken()) | |
| 94 | + monkeypatch.setattr(rl, "_sha", None) | |
| 95 | + monkeypatch.setattr(rl, "_down_until", 0.0) | |
| 96 | + monkeypatch.setattr(rl, "_last_warn", 0.0) | |
| 97 | + with caplog.at_level("WARNING"): | |
| 98 | + assert rl.apply_tier("ip:x", tiers.TIERS["keyless"], at_ms=T0) is None | |
| 99 | + assert "fails open" in caplog.text | |
| 100 | + assert not rl.available() # circuit open for a few seconds | |
| 101 | + | |
| 102 | + | |
| 103 | +def test_tier_table_is_exact(tiers): | |
| 104 | + t = tiers.TIERS | |
| 105 | + assert (t["keyless"].window_s, t["keyless"].requests, t["keyless"].rows, t["keyless"].max_rows_per_request) == (3600, 30, 100_000, 5_000) | |
| 106 | + assert (t["free"].window_s, t["free"].requests, t["free"].rows, t["free"].max_rows_per_request) == (60, 120, 1_000_000, 50_000) | |
| 107 | + assert (t["high_usage"].window_s, t["high_usage"].requests, t["high_usage"].rows, t["high_usage"].max_rows_per_request) == (60, 600, 10_000_000, 200_000) | |
| 108 | + assert t["keyless"].requests_type == "requests_per_hour" and t["free"].rows_type == "rows_per_minute" | |
| 109 | + assert tiers.row_cost(101, status=200, media_type="application/vnd.apache.parquet", quota_exempt=False) == 51 | |
| 110 | + assert tiers.row_cost(500, status=304, media_type="application/json", quota_exempt=False) == 0 | |
| 111 | + assert tiers.row_cost(500, status=200, media_type="application/json", quota_exempt=True) == 0 | |
| 112 | + assert tiers.row_cost(500, status=404, media_type="application/json", quota_exempt=False) == 0 | |
added
tests/test_ratelimit_middleware.py
+200 −0
@@ -0,0 +1,200 @@ | ||
| 1 | +"""ASGI middleware contract: headers on every response, 429 envelope, keyless vs key, costs, key gating, fail-open.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +HDRS = ["X-RateLimit-Limit-Requests", "X-RateLimit-Remaining-Requests", "X-RateLimit-Limit-Rows", | |
| 5 | + "X-RateLimit-Remaining-Rows", "X-RateLimit-Reset"] | |
| 6 | + | |
| 7 | + | |
| 8 | +def ip(n: int) -> dict: | |
| 9 | + return {"X-Forwarded-For": f"10.0.0.{n}, 172.16.0.1"} | |
| 10 | + | |
| 11 | + | |
| 12 | +def test_headers_on_success_and_error_responses(client): | |
| 13 | + r = client.get("/v1/status", headers=ip(1)) | |
| 14 | + assert r.status_code == 200 | |
| 15 | + for h in HDRS: | |
| 16 | + assert h in r.headers, h | |
| 17 | + assert r.headers["X-RateLimit-Limit-Requests"] == "30" and r.headers["X-RateLimit-Remaining-Requests"] == "29" | |
| 18 | + assert r.headers["X-RateLimit-Limit-Rows"] == "100000" | |
| 19 | + r = client.get("/v1/bars/stock/NOPE?timeframe=1day", headers=ip(1)) | |
| 20 | + assert r.status_code == 404 and r.headers["X-RateLimit-Remaining-Requests"] == "28" | |
| 21 | + assert "X-RateLimit-Reset" in r.headers | |
| 22 | + # exempt paths carry no counters and are not charged | |
| 23 | + r = client.get("/health", headers=ip(1)) | |
| 24 | + assert r.status_code == 200 and "X-RateLimit-Limit-Requests" not in r.headers | |
| 25 | + assert client.get("/v1/status", headers=ip(1)).headers["X-RateLimit-Remaining-Requests"] == "27" | |
| 26 | + | |
| 27 | + | |
| 28 | +def test_rows_are_charged_from_x_row_count(client): | |
| 29 | + r = client.get("/v1/bars/stock/AAPL?timeframe=1day&limit=100", headers=ip(2)) | |
| 30 | + assert r.headers["X-Row-Count"] == "100" | |
| 31 | + assert r.headers["X-RateLimit-Remaining-Rows"] == str(100_000 - 100) | |
| 32 | + r = client.get("/v1/_test/frame?n=1000", headers=ip(2)) | |
| 33 | + assert r.headers["X-RateLimit-Remaining-Rows"] == str(100_000 - 1100) | |
| 34 | + | |
| 35 | + | |
| 36 | +def test_parquet_costs_half_and_csv_full(client): | |
| 37 | + r = client.get("/v1/_test/frame?n=1001&format=parquet", headers=ip(3)) | |
| 38 | + assert r.status_code == 200 and r.headers["content-type"].startswith("application/vnd.apache.parquet") | |
| 39 | + assert r.headers["X-RateLimit-Remaining-Rows"] == str(100_000 - 501) | |
| 40 | + r = client.get("/v1/_test/frame?n=1001&format=csv", headers=ip(3)) | |
| 41 | + assert r.headers["X-RateLimit-Remaining-Rows"] == str(100_000 - 501 - 1001) | |
| 42 | + | |
| 43 | + | |
| 44 | +def test_quota_exempt_and_request_cost(client): | |
| 45 | + r = client.get("/v1/_test/exempt?n=5000", headers=ip(4)) | |
| 46 | + assert r.headers["X-Row-Count"] == "5000" and r.headers["X-RateLimit-Remaining-Rows"] == "100000" | |
| 47 | + assert r.headers["X-RateLimit-Remaining-Requests"] == "29" | |
| 48 | + r = client.get("/v1/_test/expensive", headers=ip(4)) | |
| 49 | + assert r.headers["X-RateLimit-Remaining-Requests"] == "27" # cost 2 | |
| 50 | + | |
| 51 | + | |
| 52 | +def test_keyless_429_envelope_retry_after_and_upgrade_hint(client): | |
| 53 | + for _ in range(30): | |
| 54 | + assert client.get("/v1/status", headers=ip(5)).status_code == 200 | |
| 55 | + r = client.get("/v1/status", headers=ip(5)) | |
| 56 | + assert r.status_code == 429 | |
| 57 | + body = r.json() | |
| 58 | + assert body["error"]["code"] == "RATE_LIMIT_EXCEEDED" and body["error"]["type"] == "requests_per_hour" | |
| 59 | + assert "free account" in body["error"]["message"] and body["detail"] | |
| 60 | + assert int(r.headers["Retry-After"]) >= 1 | |
| 61 | + assert r.headers["X-RateLimit-Remaining-Requests"] == "0" | |
| 62 | + assert body["error"]["docs"].endswith("#rate_limit_exceeded") | |
| 63 | + # another IP is unaffected | |
| 64 | + assert client.get("/v1/status", headers=ip(6)).status_code == 200 | |
| 65 | + | |
| 66 | + | |
| 67 | +def test_keyless_rows_per_hour_429(client): | |
| 68 | + for _ in range(20): | |
| 69 | + assert client.get("/v1/_test/frame?n=5000", headers=ip(7)).status_code == 200 | |
| 70 | + r = client.get("/v1/status", headers=ip(7)) | |
| 71 | + assert r.status_code == 429 and r.json()["error"]["type"] == "rows_per_hour" | |
| 72 | + assert r.headers["X-RateLimit-Remaining-Rows"] == "0" | |
| 73 | + | |
| 74 | + | |
| 75 | +def test_key_gets_free_tier_limits_and_query_param_works(client, make_user): | |
| 76 | + _, key, _, _ = make_user() | |
| 77 | + r = client.get("/v1/status", headers={**ip(8), "Authorization": f"Bearer {key}"}) | |
| 78 | + assert r.status_code == 200 | |
| 79 | + assert r.headers["X-RateLimit-Limit-Requests"] == "120" and r.headers["X-RateLimit-Limit-Rows"] == "1000000" | |
| 80 | + assert r.headers["X-RateLimit-Remaining-Requests"] == "119" | |
| 81 | + r = client.get(f"/v1/status?api_key={key}", headers=ip(8)) | |
| 82 | + assert r.headers["X-RateLimit-Remaining-Requests"] == "118" | |
| 83 | + # keyless traffic from the same IP is a different principal | |
| 84 | + assert client.get("/v1/status", headers=ip(8)).headers["X-RateLimit-Limit-Requests"] == "30" | |
| 85 | + | |
| 86 | + | |
| 87 | +def test_high_usage_tier(client, make_user): | |
| 88 | + _, key, _, _ = make_user(tier="high_usage") | |
| 89 | + r = client.get("/v1/status", headers={"Authorization": f"Bearer {key}"}) | |
| 90 | + assert r.headers["X-RateLimit-Limit-Requests"] == "600" and r.headers["X-RateLimit-Limit-Rows"] == "10000000" | |
| 91 | + | |
| 92 | + | |
| 93 | +def test_free_429_type_is_per_minute(client, make_user, monkeypatch): | |
| 94 | + _, key, _, _ = make_user() | |
| 95 | + from ratelimit import redis_limiter as rl | |
| 96 | + t0 = 1_800_000_000_000 | |
| 97 | + monkeypatch.setattr(rl, "now_ms", lambda: t0) | |
| 98 | + h = {"Authorization": f"Bearer {key}"} | |
| 99 | + for _ in range(120): | |
| 100 | + assert client.get("/v1/status", headers=h).status_code == 200 | |
| 101 | + r = client.get("/v1/status", headers=h) | |
| 102 | + assert r.status_code == 429 and r.json()["error"]["type"] == "requests_per_minute" | |
| 103 | + assert r.headers["Retry-After"] == "60" | |
| 104 | + assert "high-usage" in r.json()["error"]["message"] | |
| 105 | + monkeypatch.setattr(rl, "now_ms", lambda: t0 + 61_000) # window rolled over | |
| 106 | + assert client.get("/v1/status", headers=h).status_code == 200 | |
| 107 | + | |
| 108 | + | |
| 109 | +def test_invalid_key_is_401_not_keyless(client): | |
| 110 | + r = client.get("/v1/status", headers={"Authorization": "Bearer hfmd_live_doesnotexist000000000000000000"}) | |
| 111 | + assert r.status_code == 401 and r.json()["error"]["code"] == "INVALID_API_KEY" | |
| 112 | + r = client.get("/v1/status?api_key=garbage") | |
| 113 | + assert r.status_code == 401 | |
| 114 | + | |
| 115 | + | |
| 116 | +def test_keyless_cannot_use_key_required_endpoints(client, make_user): | |
| 117 | + r = client.get("/v1/_test/tagged", headers=ip(9)) # route tagged `stream` | |
| 118 | + assert r.status_code == 401 and r.json()["error"]["code"] == "AUTH_REQUIRED" | |
| 119 | + assert "free account" in r.json()["error"]["message"] | |
| 120 | + r = client.get("/v1/_test/needkey", headers=ip(9)) # handler sets request.state.requires_key | |
| 121 | + assert r.status_code == 401 and r.json()["error"]["code"] == "AUTH_REQUIRED" | |
| 122 | + assert "secret" not in r.text | |
| 123 | + _, key, _, _ = make_user() | |
| 124 | + assert client.get("/v1/_test/tagged", headers={"Authorization": f"Bearer {key}"}).status_code == 200 | |
| 125 | + assert client.get("/v1/_test/needkey", headers={"Authorization": f"Bearer {key}"}).json()["data"]["secret"] is True | |
| 126 | + | |
| 127 | + | |
| 128 | +def test_rows_cap_per_request(client, make_user): | |
| 129 | + r = client.get("/v1/_test/frame?n=10&limit=5001", headers=ip(10)) | |
| 130 | + assert r.status_code == 400 | |
| 131 | + body = r.json() | |
| 132 | + assert body["error"]["code"] == "ROW_LIMIT_EXCEEDED" and body["error"]["details"]["max_rows"] == 5000 | |
| 133 | + assert client.get("/v1/_test/frame?n=10&limit=5000", headers=ip(10)).status_code == 200 | |
| 134 | + _, key, _, _ = make_user() | |
| 135 | + assert client.get("/v1/_test/frame?n=10&limit=5001", headers={"Authorization": f"Bearer {key}"}).status_code == 200 | |
| 136 | + r = client.get("/v1/_test/frame?n=10&limit=50001", headers={"Authorization": f"Bearer {key}"}) | |
| 137 | + assert r.status_code == 400 and r.json()["error"]["details"]["max_rows"] == 50000 | |
| 138 | + | |
| 139 | + | |
| 140 | +def test_auth_endpoints_throttled_per_ip(client): | |
| 141 | + h = {**ip(11), "Content-Type": "application/json"} | |
| 142 | + for _ in range(10): | |
| 143 | + r = client.post("/v1/auth/login", json={"email": "nobody@example.com", "password": "x"}, headers=h) | |
| 144 | + assert r.status_code == 401 | |
| 145 | + r = client.post("/v1/auth/login", json={"email": "nobody@example.com", "password": "x"}, headers=h) | |
| 146 | + assert r.status_code == 429 | |
| 147 | + assert r.json()["error"]["type"] == "requests_per_hour" and "Retry-After" in r.headers | |
| 148 | + assert "authentication attempts" in r.json()["error"]["message"] | |
| 149 | + # the data quota of that IP is untouched | |
| 150 | + assert client.get("/v1/status", headers=ip(11)).headers["X-RateLimit-Remaining-Requests"] == "29" | |
| 151 | + | |
| 152 | + | |
| 153 | +def test_unhandled_exception_gets_envelope_and_headers(client): | |
| 154 | + r = client.get("/v1/_test/boom", headers=ip(12)) | |
| 155 | + assert r.status_code == 500 and r.json()["error"]["code"] == "INTERNAL_ERROR" | |
| 156 | + assert "X-RateLimit-Remaining-Requests" in r.headers | |
| 157 | + | |
| 158 | + | |
| 159 | +def test_fail_open_when_redis_down(client, monkeypatch): | |
| 160 | + import redis as redis_lib | |
| 161 | + | |
| 162 | + from ratelimit import redis_limiter as rl | |
| 163 | + | |
| 164 | + class Broken: | |
| 165 | + def __getattr__(self, name): | |
| 166 | + def _raise(*a, **k): | |
| 167 | + raise redis_lib.exceptions.ConnectionError("down") | |
| 168 | + return _raise | |
| 169 | + | |
| 170 | + monkeypatch.setattr(rl, "_client", Broken()) | |
| 171 | + monkeypatch.setattr(rl, "_sha", None) | |
| 172 | + monkeypatch.setattr(rl, "_down_until", 0.0) | |
| 173 | + r = client.get("/v1/status", headers=ip(13)) | |
| 174 | + assert r.status_code == 200 and "X-RateLimit-Limit-Requests" not in r.headers | |
| 175 | + r = client.get("/v1/limits", headers=ip(13)) | |
| 176 | + assert r.status_code == 200 and r.json()["data"]["principal"]["redis"] is False | |
| 177 | + | |
| 178 | + | |
| 179 | +def test_limits_endpoint_public_and_uncharged(client, make_user): | |
| 180 | + r = client.get("/v1/limits", headers=ip(14)) | |
| 181 | + assert r.status_code == 200 | |
| 182 | + d = r.json()["data"] | |
| 183 | + assert set(d["tiers"]) == {"keyless", "free", "high_usage"} | |
| 184 | + assert d["tiers"]["free"]["requests"] == 120 and d["tiers"]["keyless"]["max_rows_per_request"] == 5000 | |
| 185 | + assert d["principal"]["kind"] == "keyless" and d["principal"]["requests"]["remaining"] == 30 | |
| 186 | + assert client.get("/v1/limits", headers=ip(14)).json()["data"]["principal"]["requests"]["remaining"] == 30 | |
| 187 | + _, key, _, _ = make_user() | |
| 188 | + client.get("/v1/status", headers={"Authorization": f"Bearer {key}"}) | |
| 189 | + d = client.get("/v1/limits", headers={"Authorization": f"Bearer {key}"}).json()["data"] | |
| 190 | + assert d["principal"]["tier"] == "free" and d["principal"]["requests"]["remaining"] == 119 | |
| 191 | + | |
| 192 | + | |
| 193 | +def test_openapi_lists_new_routes_with_errors(client): | |
| 194 | + spec = client.get("/openapi.json").json() | |
| 195 | + op = spec["paths"]["/v1/auth/signup"]["post"] | |
| 196 | + assert op["summary"] and "409" in op["responses"] and "EMAIL_TAKEN" in op["responses"]["409"]["description"] | |
| 197 | + assert "429" in op["responses"] | |
| 198 | + assert "X-RateLimit-Limit-Requests" in spec["paths"]["/v1/limits"]["get"]["responses"]["200"]["headers"] | |
| 199 | + assert "EMAIL_TAKEN" in spec["components"]["schemas"]["Error"]["properties"]["error"]["properties"]["code"]["enum"] | |
| 200 | + assert spec["paths"]["/v1/me/keys"]["post"]["responses"]["201"]["content"]["application/json"]["example"]["data"]["key"].startswith("hfmd_live_") | |
added
tests/test_ratelimit_usage.py
+63 −0
@@ -0,0 +1,63 @@ | ||
| 1 | +"""Usage accounting: Redis minute counters → SQLite fold → series for the dashboard.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from datetime import datetime, timezone | |
| 5 | + | |
| 6 | + | |
| 7 | +def test_record_fold_and_series(app): | |
| 8 | + from accounts.models import UsageDaily, UsageMinute | |
| 9 | + from core.db import session | |
| 10 | + from ratelimit import redis_limiter as rl | |
| 11 | + from ratelimit import usage | |
| 12 | + | |
| 13 | + rl.reset_for_tests() | |
| 14 | + now = 1_800_000_000 # epoch s, minute-aligned | |
| 15 | + p = "key:4242" | |
| 16 | + usage.record(p, requests=1, rows=1000, status=200, ts_s=now - 180) | |
| 17 | + usage.record(p, requests=2, rows=500, rows_parquet=500, bytes_=1234, status=200, ts_s=now - 180) | |
| 18 | + usage.record(p, requests=0, status=429, ts_s=now - 120) | |
| 19 | + usage.record(p, requests=1, rows=10, status=200, ts_s=now + 5) # current minute: must stay live | |
| 20 | + assert rl.client().scard(usage.PENDING_SET) == 3 | |
| 21 | + | |
| 22 | + folded = usage.fold(now_s=now + 10) | |
| 23 | + assert folded == 2 | |
| 24 | + assert rl.client().scard(usage.PENDING_SET) == 1 # the live minute remains | |
| 25 | + with session() as s: | |
| 26 | + m = s.get(UsageMinute, (datetime.fromtimestamp(now - 180, tz=timezone.utc).replace(tzinfo=None), p)) | |
| 27 | + assert m.requests == 3 and m.rows == 1500 | |
| 28 | + d = s.get(UsageDaily, (datetime.fromtimestamp(now - 180, tz=timezone.utc).date(), p)) | |
| 29 | + assert (d.requests, d.rows, d.rows_parquet, d.bytes, d.status_2xx, d.status_429) == (3, 1500, 500, 1234, 2, 1) | |
| 30 | + # folding again is a no-op (idempotent) | |
| 31 | + assert usage.fold(now_s=now + 10) == 0 | |
| 32 | + | |
| 33 | + series = usage.usage_series([p], "24h", now_s=now + 10) | |
| 34 | + assert series["step_seconds"] == 60 and len(series["points"]) == 1440 | |
| 35 | + assert series["totals"] == {"requests": 4, "rows": 1510} # folded + live minute | |
| 36 | + by_t = {pt["t"]: pt for pt in series["points"]} | |
| 37 | + assert by_t[datetime.fromtimestamp(now - 180, tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")]["rows"] == 1500 | |
| 38 | + week = usage.usage_series([p], "7d", now_s=now + 10) | |
| 39 | + assert week["step_seconds"] == 3600 and week["totals"]["requests"] == 4 | |
| 40 | + month = usage.usage_series([p], "30d", now_s=now + 10) | |
| 41 | + assert month["step_seconds"] == 86400 and month["totals"]["rows"] == 1510 | |
| 42 | + assert usage.usage_series([], "24h", now_s=now)["totals"] == {"requests": 0, "rows": 0} | |
| 43 | + assert any(t["principal"] == p and t["requests"] >= 3 for t in usage.top_principals(days=100_000)) | |
| 44 | + | |
| 45 | + | |
| 46 | +def test_middleware_records_usage(client, make_user): | |
| 47 | + from ratelimit import usage | |
| 48 | + uid, key, _, _ = make_user() | |
| 49 | + client.get("/v1/_test/frame?n=250", headers={"Authorization": f"Bearer {key}"}) | |
| 50 | + client.get("/v1/_test/frame?n=250&format=parquet", headers={"Authorization": f"Bearer {key}"}) | |
| 51 | + from accounts import service | |
| 52 | + from core.db import session | |
| 53 | + with session() as s: | |
| 54 | + from accounts.models import User | |
| 55 | + principals = [k.principal for k in service.list_keys(s, s.get(User, uid))] | |
| 56 | + series = usage.usage_series(principals, "24h") | |
| 57 | + assert series["totals"] == {"requests": 2, "rows": 375} | |
| 58 | + | |
| 59 | + | |
| 60 | +def test_invalid_range_is_uniform_error(client, make_user): | |
| 61 | + _, key, _, _ = make_user() | |
| 62 | + r = client.get("/v1/me/usage?range=1y", headers={"Authorization": f"Bearer {key}"}) | |
| 63 | + assert r.status_code == 422 and r.json()["error"]["code"] == "VALIDATION_ERROR" | |
| 64 | ||