tick-stock-panel/backend/tests/test_kline_hot_path_memory.py

"""kline 每秒热路径测试 — _get_stock_info 走内存缓存, _attach_ext 复用签名缓存 value_map。

契约等值: 输出与旧的 DuckDB 扫描 / 逐配置读 parquet 路径一致。
"""
from __future__ import annotations

from types import SimpleNamespace

import polars as pl

from app.api import screener
from app.api.kline import _attach_ext, _get_stock_info
from app.services.ext_data import ExtConfig, ExtConfigStore, ExtField


class _FakeRepo:
    """最小 repo 替身: _get_stock_info/_load_ext_value_maps 只用到这些成员。"""

    def __init__(self, instruments: pl.DataFrame, data_dir) -> None:
        self._instruments = instruments
        self.store = SimpleNamespace(data_dir=data_dir, db=None)
        self.execute_one_calls = 0

    def get_instruments(self) -> pl.DataFrame:
        return self._instruments

    def execute_one(self, *args, **kwargs):
        self.execute_one_calls += 1
        raise AssertionError("热路径不得再走 DuckDB 扫 instruments parquet")


def _instruments_df() -> pl.DataFrame:
    return pl.DataFrame({
        "symbol": ["600000.SH", "000001.SZ", "510300.SH"],
        "name": ["浦发银行", "平安银行", "沪深300ETF"],
        "total_shares": [1.0e9, 2.0e9, None],
        "float_shares": [5.0e8, 1.0e9, None],
    })


# ---------------------------------------------------------------- _get_stock_info

def test_get_stock_info_reads_memory_cache():
    repo = _FakeRepo(_instruments_df(), None)
    assert _get_stock_info(repo, "600000.SH") == {
        "name": "浦发银行", "total_shares": 1.0e9, "float_shares": 5.0e8,
    }
    assert repo.execute_one_calls == 0


def test_get_stock_info_null_shares_become_none():
    repo = _FakeRepo(_instruments_df(), None)
    info = _get_stock_info(repo, "510300.SH")
    assert info["name"] == "沪深300ETF"
    assert info["total_shares"] is None
    assert info["float_shares"] is None


def test_get_stock_info_missing_symbol_returns_empty():
    repo = _FakeRepo(_instruments_df(), None)
    assert _get_stock_info(repo, "999999.SH") == {}


def test_get_stock_info_missing_columns_returns_empty_like_sql_error():
    df = pl.DataFrame({"symbol": ["600000.SH"], "name": ["浦发银行"]})
    repo = _FakeRepo(df, None)
    assert _get_stock_info(repo, "600000.SH") == {}


def test_get_stock_info_repo_failure_returns_empty():
    repo = SimpleNamespace()
    repo.get_instruments = lambda: (_ for _ in ()).throw(RuntimeError("boom"))
    assert _get_stock_info(repo, "600000.SH") == {}


# ---------------------------------------------------------------- _attach_ext

def test_attach_ext_empty_columns_untouched():
    resp = {"symbol": "600000.SH", "stock_info": {"name": "x"}, "rows": []}
    repo = _FakeRepo(_instruments_df(), None)
    out = _attach_ext(resp, repo, "600000.SH", None)
    assert out is resp
    assert "ext" not in out["stock_info"]
    out = _attach_ext(resp, repo, "600000.SH", "   ")
    assert "ext" not in out["stock_info"]


def test_attach_ext_invalid_spec_untouched():
    resp = {"symbol": "600000.SH", "stock_info": {"name": "x"}, "rows": []}
    repo = _FakeRepo(_instruments_df(), None)
    out = _attach_ext(resp, repo, "600000.SH", "foo_bar")  # 无 '.' 分隔
    assert out is resp
    assert "ext" not in out["stock_info"]


def test_attach_ext_maps_values_from_loader(monkeypatch):
    resp = {"symbol": "600000.SH", "stock_info": {"name": "x"}, "rows": []}
    original_info = resp["stock_info"]
    repo = _FakeRepo(_instruments_df(), None)
    monkeypatch.setattr(
        screener, "_load_ext_value_maps",
        lambda r, cols: {"cfg1__score": {"600000.SH": 88.5}, "cfg1__note": {}},
    )
    out = _attach_ext(resp, repo, "600000.SH", "cfg1.score,cfg1.note")
    assert out["stock_info"]["ext"] == {"cfg1__score": 88.5, "cfg1__note": None}
    # 原地更新 resp 是该函数契约; 但原 stock_info 字典对象不被修改 (与旧实现一致)
    assert out is resp
    assert "ext" not in original_info
    assert out["stock_info"] is not original_info


def test_attach_ext_loader_failure_yields_none_values(monkeypatch):
    resp = {"symbol": "600000.SH", "stock_info": {}, "rows": []}
    repo = _FakeRepo(_instruments_df(), None)

    def _boom(r, cols):
        raise RuntimeError("boom")

    monkeypatch.setattr(screener, "_load_ext_value_maps", _boom)
    out = _attach_ext(resp, repo, "600000.SH", "cfg1.score")
    assert out["stock_info"]["ext"] == {"cfg1__score": None}


# ------------------------------------------- _attach_ext 与真实 _load_ext_value_maps 集成

def _setup_ext_snapshot(tmp_path):
    store = ExtConfigStore(tmp_path)
    store.upsert(ExtConfig(
        id="benchx", label="基准扩展", mode="snapshot",
        fields=[ExtField(name="score", dtype="float")],
    ))
    cfg_dir = tmp_path / "ext_data" / "benchx"
    cfg_dir.mkdir(parents=True, exist_ok=True)
    pl.DataFrame({
        "symbol": [f"{600000 + i}.SH" for i in range(50)],
        "score": [float(i) for i in range(50)],
    }).write_parquet(cfg_dir / "part.parquet")
    return store


def test_attach_ext_real_loader_returns_parquet_values(tmp_path):
    _setup_ext_snapshot(tmp_path)
    repo = _FakeRepo(_instruments_df(), tmp_path)
    resp = {"symbol": "600005.SH", "stock_info": {"name": "x"}, "rows": []}
    out = _attach_ext(resp, repo, "600005.SH", "benchx.score")
    assert out["stock_info"]["ext"] == {"benchx__score": 5.0}
    assert isinstance(out["stock_info"]["ext"]["benchx__score"], float)


def test_attach_ext_real_loader_unknown_symbol_is_none(tmp_path):
    _setup_ext_snapshot(tmp_path)
    repo = _FakeRepo(_instruments_df(), tmp_path)
    resp = {"symbol": "999999.SH", "stock_info": {}, "rows": []}
    out = _attach_ext(resp, repo, "999999.SH", "benchx.score")
    assert out["stock_info"]["ext"] == {"benchx__score": None}