spb/hfmarketdata

1"""WebSocket /v1/stream: auth, subscribe, delivery of published events, resume, heartbeat, connection limit."""2from __future__ import annotations34import pytest5from starlette.websockets import WebSocketDisconnect6789@pytest.fixture10def keys(make_user):11    """Two real high_usage API keys (the rate-limit middleware validates keys before the stream handler)."""12    _, k1, _, _ = make_user(tier="high_usage")13    _, k2, _, _ = make_user(tier="high_usage")14    return k1, k2151617@pytest.fixture(autouse=True)18def _reset(app):19    from stream import broker20    broker.reset_for_tests()21    yield222324def test_keyless_is_refused_with_json_reason(client):25    with client.websocket_connect("/v1/stream") as ws:26        msg = ws.receive_json()27        assert msg["type"] == "error" and msg["code"] == "AUTH_REQUIRED" and msg["docs"].endswith("#auth_required")28        with pytest.raises(WebSocketDisconnect) as e:29            ws.receive_json()30        assert e.value.code == 4001313233def test_bad_key_is_refused(client):34    try:35        with client.websocket_connect("/v1/stream", headers={"Authorization": "Bearer not-a-key"}) as ws:36            msg = ws.receive_json()37            assert msg["code"] in ("AUTH_REQUIRED", "INVALID_API_KEY")38    except WebSocketDisconnect as e:   # refused by the rate-limit middleware before the handler39        assert e.code in (4401, 4001, 1008)404142def _key_id(raw: str) -> int:43    from sqlalchemy import select4445    from accounts.models import ApiKey46    from accounts.security import hash_key47    from core.db import session48    with session() as s:49        return int(s.scalar(select(ApiKey.id).where(ApiKey.key_hash == hash_key(raw))))505152def test_principal_is_the_key_id_not_a_hash(keys):53    """The stream must bill the same `key:<id>` principal as the HTTP limiter (a hashed pseudo-principal fed an54    orphan counter and resolved to the free tier)."""55    import hashlib5657    from core.config import settings58    from stream.auth import authenticate, principal_from_scope59    kid = _key_id(keys[0])60    assert authenticate(keys[0]) == f"key:{kid}"61    fake = "key:" + hashlib.sha256((settings.key_hash_salt + keys[0]).encode()).hexdigest()[:16]62    assert authenticate(keys[0]) != fake63    # well-formed but unknown keys are refused, whatever their shape64    assert authenticate("hfmd_live_" + "a" * 32) is None and authenticate("not-a-key") is None and authenticate(None) is None65    assert principal_from_scope({"state": {"principal": "key:7", "principal_kind": "key", "tier": "free"}}) == ("key:7", "key", "free")66    assert principal_from_scope({}) == (None, None, None)676869def test_well_formed_unknown_key_is_refused(client):70    with pytest.raises(WebSocketDisconnect) as e:71        with client.websocket_connect("/v1/stream?api_key=hfmd_live_" + "Z" * 32) as ws:72            ws.receive_json()73    assert e.value.code in (4401, 4001)747576def test_handler_refuses_without_middleware_context(keys):77    """`_resolve_principal`: middleware state wins; no state while the limiter is enabled → refuse, never guess."""78    import types7980    from core.config import settings81    from stream.routes import _resolve_principal82    ws = types.SimpleNamespace(scope={"state": {"principal": "key:42", "principal_kind": "key", "tier": "high_usage"}},83                               headers={"authorization": f"Bearer {keys[0]}"})84    assert _resolve_principal(ws, None) == ("key:42", "high_usage")85    ws.scope = {"state": {"principal": "ip:abc", "principal_kind": "keyless", "tier": "keyless"}}86    assert _resolve_principal(ws, keys[0]) == (None, "keyless")87    ws.scope = {}88    expected = (None, None) if settings.ratelimit_enabled else (f"key:{_key_id(keys[0])}", None)89    assert _resolve_principal(ws, keys[0]) == expected909192def test_subscribe_receive_and_accounting(keys, client, fundamentals_data):93    from stream import accounting, broker94    from stream.auth import authenticate95    principal = authenticate(keys[0])96    assert principal == f"key:{_key_id(keys[0])}"97    before = accounting.rows_charged(principal)98    with client.websocket_connect(f"/v1/stream?api_key={keys[0]}") as ws:99        hello = ws.receive_json()100        assert hello["type"] == "hello" and hello["heartbeat_seconds"] == 20101        ws.send_json({"action": "subscribe", "channel": "filings", "tickers": ["AAPL"], "forms": ["10-Q", "10-K"]})102        sub = ws.receive_json()103        assert sub["type"] == "subscribed" and sub["tickers"] == ["AAPL"] and sub["forms"] == ["10-K", "10-Q"]104        # events published by the ingestion (Redis) — one matching, two filtered out105        broker.publish({"type": "filing", "ticker": "MSFT", "form": "10-Q", "accn": "x1"})106        broker.publish({"type": "filing", "ticker": "AAPL", "form": "8-K", "accn": "x2"})107        seq = broker.publish({"type": "filing", "ticker": "AAPL", "form": "10-Q", "accn": "x3", "filed_date": "2024-05-03",108                              "summary": {"revenue": 90_753_000_000}})109        ev = ws.receive_json()110        assert ev["type"] == "filing" and ev["accn"] == "x3" and ev["seq"] == seq and ev["summary"]["revenue"] == 90_753_000_000111        ws.send_json({"action": "ping"})112        pong = ws.receive_json()113        assert pong["type"] == "pong" and pong["seq"] == seq114    assert accounting.rows_charged(principal) == before + 1115116117def test_real_filing_event_from_ingest(keys, client, fundamentals_data):118    from datetime import date119120    from fundamentals import ingest121    with client.websocket_connect(f"/v1/stream?api_key={keys[0]}") as ws:122        ws.receive_json()123        ws.send_json({"action": "subscribe", "channel": "filings", "tickers": "all"})124        ws.receive_json()125        ev = ingest.publish_filing_event(320193, "AAPL", "10-Q", date(2024, 5, 3), date(2024, 3, 30), "0000320193-24-000069",126                                         "aapl-20240330.htm")127        assert ev is not None128        got = ws.receive_json()129        assert got["ticker"] == "AAPL" and got["form"] == "10-Q" and got["filed_date"] == "2024-05-03"130        assert got["url"].endswith("/000032019324000069/aapl-20240330.htm")131        s = got["summary"]132        assert s["revenue"] == 90_753_000_000 and s["net_income"] == 23_636_000_000 and s["eps_diluted"] == 1.53133        assert s["total_assets"] == 337_411_000_000 and s["operating_cash_flow"] == 22_690_000_000134        assert s["yoy"]["revenue"] == pytest.approx(90_753 / 94_836 - 1, rel=1e-4)135136137def test_resume_token_replays_buffer(keys, client):138    from stream import broker139    s1 = broker.publish({"type": "filing", "ticker": "AAPL", "form": "10-K", "accn": "r1"})140    s2 = broker.publish({"type": "filing", "ticker": "AAPL", "form": "10-K", "accn": "r2"})141    s3 = broker.publish({"type": "filing", "ticker": "AAPL", "form": "10-K", "accn": "r3"})142    with client.websocket_connect(f"/v1/stream?api_key={keys[0]}") as ws:143        assert ws.receive_json()["seq"] == s3144        ws.send_json({"action": "subscribe", "channel": "filings", "tickers": ["AAPL"], "resume_token": s1})145        assert ws.receive_json()["resume_from"] == s1146        assert [ws.receive_json()["accn"] for _ in range(2)] == ["r2", "r3"]147        assert s2 < s3148    with client.websocket_connect(f"/v1/stream?api_key={keys[0]}") as ws:   # without resume: only new events149        ws.receive_json()150        ws.send_json({"action": "subscribe"})151        ws.receive_json()152        broker.publish({"type": "filing", "ticker": "AAPL", "form": "10-K", "accn": "r4"})153        assert ws.receive_json()["accn"] == "r4"154155156def test_protocol_errors(keys, client):157    with client.websocket_connect(f"/v1/stream?api_key={keys[0]}") as ws:158        ws.receive_json()159        ws.send_text("not json")160        assert ws.receive_json()["code"] == "VALIDATION_ERROR"161        ws.send_json({"action": "subscribe", "channel": "trades"})162        assert ws.receive_json()["code"] == "NOT_FOUND"163        ws.send_json({"action": "subscribe", "tickers": []})164        assert ws.receive_json()["code"] == "VALIDATION_ERROR"165        ws.send_json({"action": "dance"})166        assert ws.receive_json()["code"] == "VALIDATION_ERROR"167        ws.send_json({"action": "subscribe", "resume_token": "abc"})168        assert ws.receive_json()["code"] == "VALIDATION_ERROR"169170171def test_heartbeat(keys, client, monkeypatch):172    from stream import routes173    monkeypatch.setattr(routes, "HEARTBEAT_SECONDS", 0.3)174    with client.websocket_connect(f"/v1/stream?api_key={keys[1]}") as ws:175        ws.receive_json()176        beat = ws.receive_json()177        assert beat["type"] == "heartbeat" and "ts" in beat and "seq" in beat178179180def test_connection_limit_per_key(keys, client):181    from contextlib import ExitStack182183    from stream import routes184    with ExitStack() as stack:185        for _ in range(routes.MAX_CONNECTIONS_PER_KEY):186            ws = stack.enter_context(client.websocket_connect(f"/v1/stream?api_key={keys[0]}"))187            assert ws.receive_json()["type"] == "hello"188        extra = stack.enter_context(client.websocket_connect(f"/v1/stream?api_key={keys[0]}"))189        msg = extra.receive_json()190        assert msg["type"] == "error" and msg["code"] == "STREAM_CONNECTION_LIMIT" and msg["limit"] == 5191        with pytest.raises(WebSocketDisconnect) as e:192            extra.receive_json()193        assert e.value.code == 4029194    # slots are released on disconnect195    with client.websocket_connect(f"/v1/stream?api_key={keys[0]}") as ws:196        assert ws.receive_json()["type"] == "hello"197198199def test_phantom_connections_are_pruned(keys, client):200    """A worker killed mid-stream never releases its slot: members whose heartbeat is older than 90 s must not201    count (the old INCR/DECR counter stayed inflated for an hour and produced unjustified 4029s)."""202    import time203204    from stream import broker, routes205    principal = f"key:{_key_id(keys[1])}"206    r = broker.get_redis()207    key = routes.conn_key(principal)208    r.delete(key)209    stale = time.time() - routes.CONN_STALE_SECONDS - 5210    r.zadd(key, {f"dead-{i}": stale for i in range(routes.MAX_CONNECTIONS_PER_KEY)})       # 5 phantom slots211    r.zadd(key, {"alive": time.time()})                                                     # 1 real one elsewhere212    with client.websocket_connect(f"/v1/stream?api_key={keys[1]}") as ws:213        assert ws.receive_json()["type"] == "hello"                                         # phantoms pruned214        members = r.zrange(key, 0, -1)215        assert len(members) == 2 and "alive" in members and not any(m.startswith("dead-") for m in members)216    assert r.zrange(key, 0, -1) == ["alive"]                                                # own slot released217    r.delete(key)218219220def test_heartbeat_refreshes_connection_score(keys, client, monkeypatch):221    import time222223    from stream import broker, routes224    monkeypatch.setattr(routes, "HEARTBEAT_SECONDS", 0.2)225    monkeypatch.setattr(routes, "CONN_HEARTBEAT_SECONDS", 0.0)226    principal = f"key:{_key_id(keys[1])}"227    r = broker.get_redis()228    with client.websocket_connect(f"/v1/stream?api_key={keys[1]}") as ws:229        ws.receive_json()230        conn_id = r.zrange(routes.conn_key(principal), 0, -1)[0]231        s0 = r.zscore(routes.conn_key(principal), conn_id)232        assert ws.receive_json()["type"] == "heartbeat"233        assert ws.receive_json()["type"] == "heartbeat"234        assert r.zscore(routes.conn_key(principal), conn_id) >= s0 and time.time() - s0 < 5235236237def test_delivery_is_blocking_read_not_polling(keys, client):238    """Events published after subscribe arrive through XREAD BLOCK; events before a resume token are skipped;239    unsubscribe stops delivery (no reader), resubscribe with a token replays."""240    from stream import broker241    with client.websocket_connect(f"/v1/stream?api_key={keys[0]}") as ws:242        ws.receive_json()243        ws.send_json({"action": "subscribe", "tickers": ["NVDA"]})244        ws.receive_json()245        s1 = broker.publish({"type": "filing", "ticker": "NVDA", "form": "8-K", "accn": "b1"})246        assert ws.receive_json()["accn"] == "b1"247        ws.send_json({"action": "unsubscribe"})248        assert ws.receive_json()["type"] == "unsubscribed"249        s2 = broker.publish({"type": "filing", "ticker": "NVDA", "form": "8-K", "accn": "b2"})250        ws.send_json({"action": "ping"})251        assert ws.receive_json()["type"] == "pong"                     # nothing delivered while unsubscribed252        ws.send_json({"action": "subscribe", "tickers": ["NVDA"], "resume_token": s1})253        assert ws.receive_json()["resume_from"] == s1254        assert ws.receive_json()["accn"] == "b2" and s2 > s1255        s3 = broker.publish({"type": "filing", "ticker": "NVDA", "form": "8-K", "accn": "b3"})256        got = ws.receive_json()257        assert got["accn"] == "b3" and got["seq"] == s3258259260def test_accounting_hook_interface(app):261    from stream import accounting262    calls = []263    accounting.set_charger(lambda principal, rows: calls.append((principal, rows)))264    try:265        accounting.charge("key:abc", 1)266        assert calls == [("key:abc", 1)] and accounting.get_charger() is not None267    finally:268        accounting.set_charger(None)269270271def test_stream_info(client_hu):272    r = client_hu.get("/v1/stream/info")273    assert r.status_code == 200274    d = r.json()["data"]275    assert d["channels"] == ["filings"] and d["close_codes"]["4001"] == "AUTH_REQUIRED" and d["url"].startswith("wss://")276