easy_tdx_max/tests/unit/test_realtime_ws.py

"""RealtimeStreamHub（WebSocket 实时推送枢纽）单元测试。

全程用假行情客户端（不打真实网络、不受交易时段限制——hub 传 sessions=()），
覆盖：订阅触发轮询、多客户端并发订阅同/不同标的、退订引用计数与竞态、
无人订阅完全停止轮询、背压丢旧、以及 WS 端点端到端帧格式。
"""

from __future__ import annotations

import asyncio
from typing import Any

import pandas as pd
import pytest

from easy_tdx.web.realtime_hub import (
    MAX_WS_SYMBOLS,
    RealtimeStreamHub,
    parse_realtime_symbol,
)

pytest.importorskip("fastapi")


class FakeQuoteClient:
    """假 MAC 行情客户端：每次调用价格/量递增（绕过 feed 的 (price, vol) 去重）。"""

    def __init__(self) -> None:
        self.calls = 0
        self.requested: list[list[tuple[int, str]]] = []
        self._seq = 0

    async def get_stock_quotes(
        self, stocks: list[tuple[int, str]], fields: Any = None
    ) -> pd.DataFrame:
        self.calls += 1
        self.requested.append(list(stocks))
        self._seq += 1
        rows = []
        for market, code in stocks:
            base = 10.0 + (int(code) % 100) * 0.1
            rows.append(
                {
                    "market": int(market),
                    "code": code,
                    "close": base + self._seq * 0.01,
                    "vol": 1000.0 + self._seq,
                    "open": base,
                    "high": base + 1.0,
                    "low": base - 1.0,
                    "pre_close": base,
                    "amount": (base + self._seq * 0.01) * 1000.0,
                    "name": f"股票{code}",
                }
            )
        return pd.DataFrame(rows)


def _make_hub(**kwargs: Any) -> tuple[RealtimeStreamHub, FakeQuoteClient]:
    client = FakeQuoteClient()
    hub = RealtimeStreamHub(
        client,
        interval=0.1,  # feed 内部 clamp 下限
        sessions=(),  # 关闭交易时段过滤，测试不受运行时刻影响
        **kwargs,
    )
    return hub, client


async def _next_tick(queue: asyncio.Queue[dict[str, Any]], timeout: float = 5.0) -> dict[str, Any]:
    """取队列里的下一帧（跳过 status/ping 等控制帧）。"""
    deadline = asyncio.get_event_loop().time() + timeout
    while True:
        remaining = deadline - asyncio.get_event_loop().time()
        assert remaining > 0, "等待 tick 帧超时"
        frame = await asyncio.wait_for(queue.get(), timeout=remaining)
        if frame.get("type") == "tick":
            return frame


# ── 符号解析 ─────────────────────────────────────────────────────────────────


def test_parse_symbol() -> None:
    assert parse_realtime_symbol("SZ000001") == (0, "000001")
    assert parse_realtime_symbol("sh600519") == (1, "600519")
    assert parse_realtime_symbol("BJ920002") == (2, "920002")
    for bad in ("600519", "XX000001", "SZ00001a", "SZ0000001", ""):
        with pytest.raises(ValueError, match="标的格式"):
            parse_realtime_symbol(bad)


# ── 订阅 / 轮询生命周期 ──────────────────────────────────────────────────────


async def test_subscribe_starts_polling_and_pushes_frames() -> None:
    """订阅 → feed 开始轮询 → 队列收到 {symbol, price, ts} tick 帧。"""
    hub, client = _make_hub()
    assert not hub.polling  # 无人订阅不轮询

    cid, queue = hub.connect()
    await hub.subscribe(cid, "SZ000001")
    assert hub.polling
    assert hub.poll_symbols == ["SZ000001"]

    frame = await _next_tick(queue)
    assert frame["type"] == "tick"
    assert frame["symbol"] == "SZ000001"
    assert frame["market"] == "SZ"
    assert frame["code"] == "000001"
    assert frame["price"] > 0
    assert frame["ts"] > 0
    assert frame["name"] == "股票000001"

    # 第二帧（假客户端价格递增，绕过去重）
    frame2 = await _next_tick(queue)
    assert frame2["price"] != frame["price"]
    assert client.calls >= 2

    await hub.shutdown()


async def test_unsubscribe_last_client_stops_polling() -> None:
    """退订到无人订阅 → 轮询完全停止（节能语义）。"""
    hub, client = _make_hub()
    cid, queue = hub.connect()
    await hub.subscribe(cid, "SH600519")
    await _next_tick(queue)
    calls_before = client.calls

    await hub.unsubscribe(cid, "SH600519")
    await asyncio.sleep(0.35)  # 越过若干个 interval，确认不再发起新轮询
    assert not hub.polling
    assert hub.poll_symbols == []
    assert client.calls == calls_before

    await hub.shutdown()


# ── 多客户端并发 ─────────────────────────────────────────────────────────────


async def test_multiple_clients_same_symbol_all_receive() -> None:
    """两个客户端订阅同一标的：都收到推送；退订一个后轮询继续（引用计数）。"""
    hub, client = _make_hub()
    cid1, q1 = hub.connect()
    cid2, q2 = hub.connect()
    await hub.subscribe(cid1, "SZ000001")
    await hub.subscribe(cid2, "sz000001")  # 大小写归一到同一标的

    f1 = await _next_tick(q1)
    f2 = await _next_tick(q2)
    assert f1["symbol"] == f2["symbol"] == "SZ000001"
    # 去重后只轮询一份
    assert hub.poll_symbols == ["SZ000001"]

    # 客户端 1 断开：客户端 2 仍在订阅，轮询不停止
    await hub.disconnect(cid1)
    await asyncio.sleep(0.2)
    assert hub.polling
    await _next_tick(q2)

    await hub.shutdown()


async def test_clients_different_symbols_no_cross_delivery() -> None:
    """不同客户端订阅不同标的：各自只收到自己标的的帧。"""
    hub, _client = _make_hub()
    cid_a, q_a = hub.connect()
    cid_b, q_b = hub.connect()
    await hub.subscribe(cid_a, "SZ000001")
    await hub.subscribe(cid_b, "SH600519")
    assert hub.poll_symbols == ["SH600519", "SZ000001"]  # 排序确定

    for _ in range(3):  # 多收几帧确认无串扰
        fa = await _next_tick(q_a)
        fb = await _next_tick(q_b)
        assert fa["symbol"] == "SZ000001"
        assert fb["symbol"] == "SH600519"

    await hub.shutdown()


async def test_dynamic_subscribe_extends_poll_set() -> None:
    """运行中追加订阅：新标的进入轮询集合并开始推送。"""
    hub, client = _make_hub()
    cid, queue = hub.connect()
    await hub.subscribe(cid, "SZ000001")
    await _next_tick(queue)

    await hub.subscribe(cid, "SH600519")
    assert hub.poll_symbols == ["SH600519", "SZ000001"]
    seen: set[str] = set()
    for _ in range(8):
        frame = await _next_tick(queue)
        seen.add(frame["symbol"])
        if seen == {"SZ000001", "SH600519"}:
            break
    assert seen == {"SZ000001", "SH600519"}
    # 每轮轮询带全量订阅集合（requested 记录 (market_int, code) 元组）
    requested_keys = {f"SZ{c}" if m == 0 else f"SH{c}" for m, c in client.requested[-1]}
    assert {"SZ000001", "SH600519"}.issubset(requested_keys)

    await hub.shutdown()


# ── 竞态 ─────────────────────────────────────────────────────────────────────


async def test_concurrent_subscribe_unsubscribe_race() -> None:
    """多客户端并发订阅/退订同一批标的：最终引用计数与轮询状态一致。"""
    hub, _client = _make_hub()
    cids = [hub.connect()[0] for _ in range(8)]
    symbols = ["SZ000001", "SH600519", "SZ399006"]

    # 并发混订：偶数客户端先订后退，奇数只订
    await asyncio.gather(
        *[hub.subscribe(cid, sym) for cid in cids for sym in symbols],
        *[hub.unsubscribe(cid, sym) for cid in cids[::2] for sym in symbols],
    )
    # 奇数客户端（索引 1,3,5,7）仍各持有 3 个标的的订阅
    assert hub.poll_symbols == sorted(symbols)
    assert hub.subscriber_count == 8

    # 全部断开（并发）→ 轮询停止
    await asyncio.gather(*[hub.disconnect(cid) for cid in cids])
    await asyncio.sleep(0.2)
    assert not hub.polling
    assert hub.poll_symbols == []

    await hub.shutdown()


async def test_symbol_cap_rejected() -> None:
    """去重后标的数超 80（协议上限）拒绝并抛 ValueError。"""
    hub, _client = _make_hub()
    cid, _queue = hub.connect()
    for i in range(MAX_WS_SYMBOLS):
        await hub.subscribe(cid, f"SH{i:06d}")
    with pytest.raises(ValueError, match="上限"):
        await hub.subscribe(cid, "SZ000001")
    await hub.shutdown()


async def test_backpressure_drops_oldest() -> None:
    """队列满时丢最旧保最新（行情快照旧帧无价值）。"""
    hub, _client = _make_hub(queue_size=2)
    cid, queue = hub.connect()
    await hub.subscribe(cid, "SZ000001")

    # 不消费，等队列塞满并溢出
    await asyncio.sleep(0.8)
    assert queue.qsize() <= 2
    prices = [queue.get_nowait()["price"] for _ in range(queue.qsize())]
    # 留下的是最新帧（假客户端价格单调递增）
    assert prices == sorted(prices)

    await hub.shutdown()


# ── WS 端点端到端（TestClient）───────────────────────────────────────────────


def test_ws_endpoint_end_to_end(monkeypatch: pytest.MonkeyPatch) -> None:
    """连接 → 收 tick 帧（真实路由 + hub + feed，仅行情源为假）。

    hub 的 feed 任务跑在 TestClient 的 portal 事件循环里，客户端退出时随循环
    一并销毁，无需（也无法）在测试协程里再 shutdown。
    """
    from fastapi import FastAPI
    from fastapi.testclient import TestClient

    from easy_tdx.web.routers.realtime import router

    app = FastAPI()
    app.include_router(router, prefix="/api/v1")
    hub, _client = _make_hub()
    app.state.realtime_hub = hub

    with TestClient(app) as client, client.websocket_connect("/api/v1/ws/realtime/SZ000001") as ws:
        frame = ws.receive_json()
        assert frame["type"] == "tick"
        assert frame["symbol"] == "SZ000001"
        assert {"price", "ts", "market", "code"} <= set(frame)