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

"""`/rows` 查询增强 (V-开放能力): filter 等值过滤 / sort 排序 / offset 分页 /
timeseries 日期范围合并读取。

范围读取是路径敏感操作 (拼 `date=` 分区目录), 与 date_guard 同一套防穿越
校验 (_partition_date); 过滤/排序的 400 语义单独覆盖。
"""
from __future__ import annotations

from pathlib import Path

import polars as pl
import pytest
from fastapi import HTTPException

from app.api.ext_data import _apply_row_filters, _apply_row_sort, _read_ext_dataframe
from app.services.ext_data import ExtConfig, ExtField


def _cfg(mode: str = "timeseries") -> ExtConfig:
    return ExtConfig(
        id="hot",
        label="人气",
        mode=mode,  # type: ignore[arg-type]
        fields=[ExtField("symbol", "string"), ExtField("concept", "string"), ExtField("heat", "float")],
    )


def _layout(tmp_path: Path) -> Path:
    """三个时序分区: 09-09 / 09-10 / 09-11。"""
    data_dir = tmp_path / "data"
    rows_by_day = {
        "2026-09-09": [("000001.SZ", "AI", 1.0), ("600519.SH", "白酒", 2.0)],
        "2026-09-10": [("000001.SZ", "AI", 3.0)],
        "2026-09-11": [("300750.SZ", "AI", 5.0), ("600519.SH", "白酒", 4.0)],
    }
    for day, rows in rows_by_day.items():
        part = data_dir / "ext_data" / "hot" / "timeseries" / f"date={day}"
        part.mkdir(parents=True)
        pl.DataFrame(
            {"symbol": [r[0] for r in rows], "concept": [r[1] for r in rows], "heat": [r[2] for r in rows]}
        ).write_parquet(part / "part.parquet")
    return data_dir


# ── 日期范围读取 ──────────────────────────────────────────
def test_read_range_merges_partitions(tmp_path: Path) -> None:
    df, active = _read_ext_dataframe(_cfg(), _layout(tmp_path), start_date="2026-09-09", end_date="2026-09-10")
    assert len(df) == 3
    assert active == "2026-09-09..2026-09-10"
    # 范围合并自动补行级分区日期 (表本身无 date 列 → 补 date 列)
    assert "date" in df.columns
    assert set(df.filter(pl.col("symbol") == "600519.SH")["date"].to_list()) == {"2026-09-09"}


def test_read_range_open_ended(tmp_path: Path) -> None:
    data_dir = _layout(tmp_path)
    df, active = _read_ext_dataframe(_cfg(), data_dir, start_date="2026-09-10")
    assert len(df) == 3 and active == "2026-09-10..2026-09-11"
    df, active = _read_ext_dataframe(_cfg(), data_dir, end_date="2026-09-09")
    assert len(df) == 2 and active == "2026-09-09..2026-09-09"


def test_read_range_reversed_rejected(tmp_path: Path) -> None:
    with pytest.raises(HTTPException) as e:
        _read_ext_dataframe(_cfg(), _layout(tmp_path), start_date="2026-09-11", end_date="2026-09-09")
    assert e.value.status_code == 400


def test_read_range_bad_date_rejected(tmp_path: Path) -> None:
    with pytest.raises(HTTPException) as e:
        _read_ext_dataframe(_cfg(), _layout(tmp_path), start_date="../../kline_daily")
    assert e.value.status_code == 400


def test_snapshot_mode_rejects_range(tmp_path: Path) -> None:
    data_dir = _layout(tmp_path)
    snap = data_dir / "ext_data" / "hot" / "part.parquet"
    snap.parent.mkdir(parents=True, exist_ok=True)
    pl.DataFrame({"symbol": ["000001.SZ"], "concept": ["AI"], "heat": [1.0]}).write_parquet(snap)
    with pytest.raises(HTTPException) as e:
        _read_ext_dataframe(_cfg(mode="snapshot"), data_dir, start_date="2026-09-09")
    assert e.value.status_code == 400


def test_exact_date_still_wins_over_range(tmp_path: Path) -> None:
    df, active = _read_ext_dataframe(
        _cfg(), _layout(tmp_path),
        snapshot_date="2026-09-10", start_date="2026-09-09", end_date="2026-09-11",
    )
    assert len(df) == 1 and active == "2026-09-10"


def test_read_range_tolerates_schema_drift(tmp_path: Path) -> None:
    """分区 schema 不一致 (配置中途加字段 → 早期分区少列): diagonal 合并, 缺列补 null。

    真实场景: ext_fuyao_hot 早期 4 列、后期 7 列, vertical concat 直接 ShapeError。
    """
    data_dir = _layout(tmp_path)
    # 给最后一天追加一列新字段 (模拟中途加列)
    last = data_dir / "ext_data" / "hot" / "timeseries" / "date=2026-09-11" / "part.parquet"
    df = pl.read_parquet(last).with_columns(pl.lit("x").alias("extra_tag"))
    df.write_parquet(last)

    out, _active = _read_ext_dataframe(_cfg(), data_dir, start_date="2026-09-09", end_date="2026-09-11")
    assert len(out) == 5
    assert "extra_tag" in out.columns
    # 早期分区的 extra_tag 为 null, 后期有值
    tags = out.sort("symbol").get_column("extra_tag").to_list()
    assert None in tags and "x" in tags


# ── 过滤 / 排序 ────────────────────────────────────────────
def _df() -> pl.DataFrame:
    return pl.DataFrame({
        "symbol": ["000001.SZ", "600519.SH", "300750.SZ"],
        "concept": ["AI", "白酒", "AI"],
        "heat": [1.0, 2.5, 3.0],
    })


def test_filter_single_value_and_in_list() -> None:
    out = _apply_row_filters(_df(), ["concept:AI"])
    assert out["symbol"].to_list() == ["000001.SZ", "300750.SZ"]
    # 值间 OR
    out = _apply_row_filters(_df(), ["concept:AI|白酒"])
    assert len(out) == 3
    # 数值列按字符串等值匹配
    out = _apply_row_filters(_df(), ["heat:2.5"])
    assert out["symbol"].to_list() == ["600519.SH"]


def test_filter_multiple_conditions_are_and() -> None:
    out = _apply_row_filters(_df(), ["concept:AI", "heat:3.0"])
    assert out["symbol"].to_list() == ["300750.SZ"]


def test_filter_unknown_field_and_malformed_rejected() -> None:
    with pytest.raises(HTTPException) as e:
        _apply_row_filters(_df(), ["nope:1"])
    assert e.value.status_code == 400
    with pytest.raises(HTTPException):
        _apply_row_filters(_df(), ["concept"])  # 缺 :值
    with pytest.raises(HTTPException):
        _apply_row_filters(_df(), [":AI"])  # 缺字段


def test_sort_asc_desc_and_unknown_rejected() -> None:
    out = _apply_row_sort(_df(), "heat")
    assert out["heat"].to_list() == [1.0, 2.5, 3.0]
    out = _apply_row_sort(_df(), "heat:desc")
    assert out["heat"].to_list() == [3.0, 2.5, 1.0]
    with pytest.raises(HTTPException) as e:
        _apply_row_sort(_df(), "nope")
    assert e.value.status_code == 400