spb/hfmarketdata

1"""Usage accounting: Redis minute counters → SQLite fold → series for the dashboard."""2from __future__ import annotations34from datetime import datetime, timezone567def test_record_fold_and_series(app):8    from accounts.models import UsageDaily, UsageMinute9    from core.db import session10    from ratelimit import redis_limiter as rl11    from ratelimit import usage1213    rl.reset_for_tests()14    now = 1_800_000_000          # epoch s, minute-aligned15    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 live20    assert rl.client().scard(usage.PENDING_SET) == 32122    folded = usage.fold(now_s=now + 10)23    assert folded == 224    assert rl.client().scard(usage.PENDING_SET) == 1          # the live minute remains25    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 == 150028        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) == 03233    series = usage.usage_series([p], "24h", now_s=now + 10)34    assert series["step_seconds"] == 60 and len(series["points"]) == 144035    assert series["totals"] == {"requests": 4, "rows": 1510, "status_429": 1, "bytes": 1234, "rows_parquet": 500}   # folded + live minute36    assert m.status_429 == 0 and m.rows_parquet == 500 and m.bytes == 1234   # usage_minute carries the new columns (the 429 is in the next minute)37    by_t = {pt["t"]: pt for pt in series["points"]}38    assert by_t[datetime.fromtimestamp(now - 180, tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")]["rows"] == 150039    week = usage.usage_series([p], "7d", now_s=now + 10)40    assert week["step_seconds"] == 3600 and week["totals"]["requests"] == 441    month = usage.usage_series([p], "30d", now_s=now + 10)42    assert month["step_seconds"] == 86400 and month["totals"]["rows"] == 151043    assert usage.usage_series([], "24h", now_s=now)["totals"] == {"requests": 0, "rows": 0, "status_429": 0, "bytes": 0, "rows_parquet": 0}44    csv = usage.usage_csv(series).splitlines()45    assert csv[0] == "t,requests,rows,status_429,bytes,rows_parquet" and len(csv) == 144146    assert any(line.endswith(",3,1500,0,1234,500") for line in csv[1:])47    assert any(t["principal"] == p and t["requests"] >= 3 for t in usage.top_principals(days=100_000))484950def test_middleware_records_usage(client, make_user):51    from ratelimit import usage52    uid, key, _, _ = make_user()53    client.get("/v1/_test/frame?n=250", headers={"Authorization": f"Bearer {key}"})54    client.get("/v1/_test/frame?n=250&format=parquet", headers={"Authorization": f"Bearer {key}"})55    from accounts import service56    from core.db import session57    with session() as s:58        from accounts.models import User59        principals = [k.principal for k in service.list_keys(s, s.get(User, uid))]60    series = usage.usage_series(principals, "24h")61    assert (series["totals"]["requests"], series["totals"]["rows"], series["totals"]["rows_parquet"]) == (2, 375, 125)62    assert series["totals"]["bytes"] > 0 and series["totals"]["status_429"] == 0636465def test_invalid_range_is_uniform_error(client, make_user):66    _, key, _, _ = make_user()67    r = client.get("/v1/me/usage?range=1y", headers={"Authorization": f"Bearer {key}"})68    assert r.status_code == 422 and r.json()["error"]["code"] == "VALIDATION_ERROR"697071def test_quota_alerts_once_per_day_and_opt_out(app, make_user, outbox):72    """B4: 80 % / 100 % of the rows window or a 429 → one mail per user per day; opt-out honoured."""73    from accounts.models import User74    from core.db import session75    from ratelimit import redis_limiter as rl76    from ratelimit import usage77    from ratelimit.tiers import TIERS78    uid, key, email, _ = make_user()79    with session() as s:80        from accounts import service81        principal = service.active_keys(s, s.get(User, uid))[0].principal82    rl.reset_for_tests()83    free = TIERS["free"]84    rl.apply_tier(principal, free, req_cost=1, rows_cost=int(free.rows * 0.85), force=True)   # 85 % of the window85    usage.record(principal, requests=1, rows=10, status=200)                                     # "recent" principal86    assert usage.quota_alerts() == 187    alert = [m for m in outbox if m.kind == "quota_alert"][-1]88    assert alert.to == email and "80 %" in alert.ctx["what"]89    assert usage.quota_alerts() == 0                                                              # deduped for today90    with session() as s:91        s.get(User, uid).quota_alerts = 092    rl._call(lambda: rl.client().delete(*[k for k in rl.client().keys("alert:*")]) if rl.client().keys("alert:*") else None)93    assert usage.quota_alerts() == 0                                                              # opt-out94    with session() as s:95        s.get(User, uid).quota_alerts = 196    usage.record(principal, requests=0, status=429)97    assert usage.quota_alerts() == 198    alert = [m for m in outbox if m.kind == "quota_alert"][-1]99    assert "100 %" in alert.ctx["what"] or "429" in alert.ctx["what"] or "80 %" in alert.ctx["what"]100