spb/hfmarketdata

1"""ASGI middleware contract: headers on every response, 429 envelope, keyless vs key, costs, key gating, fail-open."""2from __future__ import annotations34HDRS = ["X-RateLimit-Limit-Requests", "X-RateLimit-Remaining-Requests", "X-RateLimit-Limit-Rows",5        "X-RateLimit-Remaining-Rows", "X-RateLimit-Reset"]678def ip(n: int) -> dict:9    """The proxy (ngrok) APPENDS the real peer: the last hop is the client, the first one is whatever it claimed."""10    return {"X-Forwarded-For": f"172.16.0.1, 10.0.0.{n}"}111213def test_headers_on_success_and_error_responses(client):14    r = client.get("/v1/status", headers=ip(1))15    assert r.status_code == 20016    for h in HDRS:17        assert h in r.headers, h18    assert r.headers["X-RateLimit-Limit-Requests"] == "30" and r.headers["X-RateLimit-Remaining-Requests"] == "29"19    assert r.headers["X-RateLimit-Limit-Rows"] == "100000"20    r = client.get("/v1/bars/stock/NOPE?timeframe=1day", headers=ip(1))21    assert r.status_code == 404 and r.headers["X-RateLimit-Remaining-Requests"] == "28"22    assert "X-RateLimit-Reset" in r.headers23    # exempt paths carry no counters and are not charged24    r = client.get("/health", headers=ip(1))25    assert r.status_code == 200 and "X-RateLimit-Limit-Requests" not in r.headers26    assert client.get("/v1/status", headers=ip(1)).headers["X-RateLimit-Remaining-Requests"] == "27"272829def test_rows_are_charged_from_x_row_count(client):30    r = client.get("/v1/bars/stock/AAPL?timeframe=1day&limit=100", headers=ip(2))31    assert r.headers["X-Row-Count"] == "100"32    assert r.headers["X-RateLimit-Remaining-Rows"] == str(100_000 - 100)33    r = client.get("/v1/_test/frame?n=1000", headers=ip(2))34    assert r.headers["X-RateLimit-Remaining-Rows"] == str(100_000 - 1100)353637def test_parquet_costs_half_and_csv_full(client):38    r = client.get("/v1/_test/frame?n=1001&format=parquet", headers=ip(3))39    assert r.status_code == 200 and r.headers["content-type"].startswith("application/vnd.apache.parquet")40    assert r.headers["X-RateLimit-Remaining-Rows"] == str(100_000 - 501)41    r = client.get("/v1/_test/frame?n=1001&format=csv", headers=ip(3))42    assert r.headers["X-RateLimit-Remaining-Rows"] == str(100_000 - 501 - 1001)434445def test_quota_exempt_and_request_cost(client):46    r = client.get("/v1/_test/exempt?n=5000", headers=ip(4))47    assert r.headers["X-Row-Count"] == "5000" and r.headers["X-RateLimit-Remaining-Rows"] == "100000"48    assert r.headers["X-RateLimit-Remaining-Requests"] == "29"49    r = client.get("/v1/_test/expensive", headers=ip(4))50    assert r.headers["X-RateLimit-Remaining-Requests"] == "27"   # cost 2515253def test_keyless_429_envelope_retry_after_and_upgrade_hint(client):54    for _ in range(30):55        assert client.get("/v1/status", headers=ip(5)).status_code == 20056    r = client.get("/v1/status", headers=ip(5))57    assert r.status_code == 42958    body = r.json()59    assert body["error"]["code"] == "RATE_LIMIT_EXCEEDED" and body["error"]["type"] == "requests_per_hour"60    assert "free account" in body["error"]["message"] and body["detail"]61    assert int(r.headers["Retry-After"]) >= 162    assert r.headers["X-RateLimit-Remaining-Requests"] == "0"63    assert body["error"]["docs"].endswith("#rate_limit_exceeded")64    # another IP is unaffected65    assert client.get("/v1/status", headers=ip(6)).status_code == 200666768def test_keyless_rows_per_hour_429(client):69    for _ in range(20):70        assert client.get("/v1/_test/frame?n=5000", headers=ip(7)).status_code == 20071    r = client.get("/v1/status", headers=ip(7))72    assert r.status_code == 429 and r.json()["error"]["type"] == "rows_per_hour"73    assert r.headers["X-RateLimit-Remaining-Rows"] == "0"747576def test_key_gets_free_tier_limits_and_query_param_works(client, make_user):77    _, key, _, _ = make_user()78    r = client.get("/v1/status", headers={**ip(8), "Authorization": f"Bearer {key}"})79    assert r.status_code == 20080    assert r.headers["X-RateLimit-Limit-Requests"] == "120" and r.headers["X-RateLimit-Limit-Rows"] == "1000000"81    assert r.headers["X-RateLimit-Remaining-Requests"] == "119"82    r = client.get(f"/v1/status?api_key={key}", headers=ip(8))83    assert r.headers["X-RateLimit-Remaining-Requests"] == "118"84    # keyless traffic from the same IP is a different principal85    assert client.get("/v1/status", headers=ip(8)).headers["X-RateLimit-Limit-Requests"] == "30"868788def test_high_usage_tier(client, make_user):89    _, key, _, _ = make_user(tier="high_usage")90    r = client.get("/v1/status", headers={"Authorization": f"Bearer {key}"})91    assert r.headers["X-RateLimit-Limit-Requests"] == "600" and r.headers["X-RateLimit-Limit-Rows"] == "10000000"929394def test_free_429_type_is_per_minute(client, make_user, monkeypatch):95    _, key, _, _ = make_user()96    from ratelimit import redis_limiter as rl97    t0 = 1_800_000_000_00098    monkeypatch.setattr(rl, "now_ms", lambda: t0)99    h = {"Authorization": f"Bearer {key}"}100    for _ in range(120):101        assert client.get("/v1/status", headers=h).status_code == 200102    r = client.get("/v1/status", headers=h)103    assert r.status_code == 429 and r.json()["error"]["type"] == "requests_per_minute"104    assert r.headers["Retry-After"] == "60"105    assert "high-usage" in r.json()["error"]["message"]106    monkeypatch.setattr(rl, "now_ms", lambda: t0 + 61_000)      # window rolled over107    assert client.get("/v1/status", headers=h).status_code == 200108109110def test_invalid_key_is_401_not_keyless(client):111    r = client.get("/v1/status", headers={"Authorization": "Bearer hfmd_live_doesnotexist000000000000000000"})112    assert r.status_code == 401 and r.json()["error"]["code"] == "INVALID_API_KEY"113    r = client.get("/v1/status?api_key=garbage")114    assert r.status_code == 401115116117def test_keyless_cannot_use_key_required_endpoints(client, make_user):118    r = client.get("/v1/_test/tagged", headers=ip(9))          # route tagged `stream`119    assert r.status_code == 401 and r.json()["error"]["code"] == "AUTH_REQUIRED"120    assert "free account" in r.json()["error"]["message"]121    r = client.get("/v1/_test/needkey", headers=ip(9))         # handler sets request.state.requires_key122    assert r.status_code == 401 and r.json()["error"]["code"] == "AUTH_REQUIRED"123    assert "secret" not in r.text124    _, key, _, _ = make_user()125    assert client.get("/v1/_test/tagged", headers={"Authorization": f"Bearer {key}"}).status_code == 200126    assert client.get("/v1/_test/needkey", headers={"Authorization": f"Bearer {key}"}).json()["data"]["secret"] is True127128129def test_rows_cap_per_request(client, make_user):130    r = client.get("/v1/_test/frame?n=10&limit=5001", headers=ip(10))131    assert r.status_code == 400132    body = r.json()133    assert body["error"]["code"] == "ROW_LIMIT_EXCEEDED" and body["error"]["details"]["max_rows"] == 5000134    assert client.get("/v1/_test/frame?n=10&limit=5000", headers=ip(10)).status_code == 200135    _, key, _, _ = make_user()136    assert client.get("/v1/_test/frame?n=10&limit=5001", headers={"Authorization": f"Bearer {key}"}).status_code == 200137    r = client.get("/v1/_test/frame?n=10&limit=50001", headers={"Authorization": f"Bearer {key}"})138    assert r.status_code == 400 and r.json()["error"]["details"]["max_rows"] == 50000139140141def test_auth_endpoints_throttled_per_ip(client):142    h = {**ip(11), "Content-Type": "application/json"}143    for _ in range(20):144        r = client.post("/v1/auth/login", json={"email": "nobody@example.com", "password": "x"}, headers=h)145        assert r.status_code == 401146    r = client.post("/v1/auth/login", json={"email": "nobody@example.com", "password": "x"}, headers=h)147    assert r.status_code == 429148    assert r.json()["error"]["type"] == "requests_per_hour" and "Retry-After" in r.headers149    assert "authentication attempts" in r.json()["error"]["message"]150    # the data quota of that IP is untouched151    assert client.get("/v1/status", headers=ip(11)).headers["X-RateLimit-Remaining-Requests"] == "29"152153154def test_unhandled_exception_gets_envelope_and_headers(client):155    r = client.get("/v1/_test/boom", headers=ip(12))156    assert r.status_code == 500 and r.json()["error"]["code"] == "INTERNAL_ERROR"157    assert "X-RateLimit-Remaining-Requests" in r.headers158159160def test_fail_open_when_redis_down(client, monkeypatch):161    import redis as redis_lib162163    from ratelimit import redis_limiter as rl164165    class Broken:166        def __getattr__(self, name):167            def _raise(*a, **k):168                raise redis_lib.exceptions.ConnectionError("down")169            return _raise170171    monkeypatch.setattr(rl, "_client", Broken())172    monkeypatch.setattr(rl, "_sha", None)173    monkeypatch.setattr(rl, "_down_until", 0.0)174    r = client.get("/v1/status", headers=ip(13))175    assert r.status_code == 200 and "X-RateLimit-Limit-Requests" not in r.headers176    r = client.get("/v1/limits", headers=ip(13))177    assert r.status_code == 200 and r.json()["data"]["principal"]["redis"] is False178179180def test_limits_endpoint_public_and_uncharged(client, make_user):181    r = client.get("/v1/limits", headers=ip(14))182    assert r.status_code == 200183    d = r.json()["data"]184    assert set(d["tiers"]) == {"keyless", "free", "high_usage"}185    assert d["tiers"]["free"]["requests"] == 120 and d["tiers"]["keyless"]["max_rows_per_request"] == 5000186    assert d["principal"]["kind"] == "keyless" and d["principal"]["requests"]["remaining"] == 30187    assert client.get("/v1/limits", headers=ip(14)).json()["data"]["principal"]["requests"]["remaining"] == 30188    _, key, _, _ = make_user()189    client.get("/v1/status", headers={"Authorization": f"Bearer {key}"})190    d = client.get("/v1/limits", headers={"Authorization": f"Bearer {key}"}).json()["data"]191    assert d["principal"]["tier"] == "free" and d["principal"]["requests"]["remaining"] == 119192193194def test_openapi_lists_new_routes_with_errors(client):195    spec = client.get("/openapi.json").json()196    op = spec["paths"]["/v1/auth/signup"]["post"]197    assert op["summary"] and "202" in op["responses"] and "400" in op["responses"] and "WEAK_PASSWORD" in op["responses"]["400"]["description"]198    assert "429" in op["responses"]199    op = spec["paths"]["/v1/admin/users"]["post"]200    assert "409" in op["responses"] and "EMAIL_TAKEN" in op["responses"]["409"]["description"]201    assert "X-RateLimit-Limit-Requests" in spec["paths"]["/v1/limits"]["get"]["responses"]["200"]["headers"]202    assert "EMAIL_TAKEN" in spec["components"]["schemas"]["Error"]["properties"]["error"]["properties"]["code"]["enum"]203    assert spec["paths"]["/v1/me/keys"]["post"]["responses"]["201"]["content"]["application/json"]["example"]["data"]["key"].startswith("hfmd_live_")204205206def test_browser_session_gets_account_tier_on_data_endpoints(web, make_user, signin):207    """A signed-in browser (cookie, no Bearer) is the `user:<id>` principal with the account's tier — the charts208    page and the playground get 120 req/min without pasting a key. Anonymous cookies stay keyless."""209    uid, _raw, email, pw = make_user(tier="free", with_key=False)210    r = web.get("/v1/status", headers=ip(41))211    assert r.headers["X-RateLimit-Limit-Requests"] == "30"                      # keyless before sign-in212    signin(web, email, pw)213    r = web.get("/v1/bars/stock/AAPL?timeframe=1day&limit=5", headers=ip(41))214    assert r.status_code == 200215    assert r.headers["X-RateLimit-Limit-Requests"] == "120" and r.headers["X-RateLimit-Limit-Rows"] == "1000000"216    lim = web.get("/v1/limits", headers=ip(41)).json()["data"]["principal"]217    assert lim["principal"] == f"user:{uid}" and lim["kind"] == "session" and lim["tier"] == "free"218    # the session's usage is part of the account's dashboard series219    me = web.get("/v1/me/usage?range=24h", headers={"X-Requested-With": "hfmd"}).json()["data"]220    assert f"user:{uid}" in me["principals"]221    # a Bearer key still wins over the cookie222    uid2, raw2, _, _ = make_user(tier="high_usage")223    r = web.get("/v1/status", headers={**ip(41), "Authorization": f"Bearer {raw2}"})224    assert r.headers["X-RateLimit-Limit-Requests"] == "600"225    # sign out everywhere → the OLD cookie (another device) is stale → keyless again; the current browser is refreshed226    old_cookie = web.cookies.get("hfmd_session")227    r = web.post("/v1/me/sessions/revoke-all", json={}, headers={"Content-Type": "application/json", "X-Requested-With": "hfmd"})228    assert r.status_code in (200, 204), r.text229    from fastapi.testclient import TestClient230    with TestClient(web.app) as other:231        r = other.get("/v1/status", headers={**ip(41), "Cookie": f"hfmd_session={old_cookie}"})232        assert r.headers["X-RateLimit-Limit-Requests"] == "30"233        assert other.get("/v1/me", headers={"Cookie": f"hfmd_session={old_cookie}"}).status_code == 401234    assert web.get("/v1/status", headers=ip(41)).headers["X-RateLimit-Limit-Requests"] == "120"235