ggx/tests/test_units.py

"""量价单位归一化（``stock_daily`` 的历史遗留混用）测试。

背景：``stock_daily`` 里 2015-01~2019 的行是 Tushare 原始单位（手 / 千元），
2020 起沿用既有 qlib 存量（股 / 元），2019 年同日混着两种。
按「元」配置的流动性阈值因此把早年低估 1000 倍，会把股票池整体清空。
这些测试锁定「读取层必须幂等地归一化」这一契约。
"""

from __future__ import annotations

import pandas as pd
import pytest

from hdiv.data.units import (
    OHLCV_CONVERTED,
    OHLCV_UNKNOWN,
    amount_qian_to_yuan,
    detect_ohlcv_units,
    normalize_ohlcv_units,
    ohlcv_unit_ratio,
    vol_shou_to_shares,
)


def _row(symbol: str, close: float, shares: float, **kw) -> dict:
    """按「股 / 元」口径构造一行（即换算后的目标形态）。"""
    return {
        "symbol": symbol,
        "close": close,
        "volume": shares,
        "amount": shares * close,
        **kw,
    }


def _raw_row(symbol: str, close: float, shares: float, **kw) -> dict:
    """按 Tushare 原始口径构造一行：成交量为手、成交额为千元。

    ``shares`` 是真实股数。手 = 股/100；千元 = (股 × 价)/1000。
    """
    return {
        "symbol": symbol,
        "close": close,
        "volume": shares / 100.0,
        "amount": shares * close / 1000.0,
        **kw,
    }


class TestScalarConversions:
    def test_vol_shou_to_shares(self) -> None:
        assert vol_shou_to_shares(pd.Series([100.0])).iloc[0] == 10000.0

    def test_amount_qian_to_yuan(self) -> None:
        assert amount_qian_to_yuan(pd.Series([1000.0])).iloc[0] == 1_000_000.0


class TestUnitRatio:
    def test_converted_rows_ratio_is_one(self) -> None:
        df = pd.DataFrame([_row("000001.SZ", 10.0, 1e6)])
        assert ohlcv_unit_ratio(df).iloc[0] == pytest.approx(1.0)

    def test_raw_rows_ratio_is_point_one(self) -> None:
        # 手 / 千元：volume 是股数/100，amount 是元/1000 → 比值 1/10
        df = pd.DataFrame([_raw_row("000001.SZ", 10.0, 1e6)])
        assert ohlcv_unit_ratio(df).iloc[0] == pytest.approx(0.1)

    def test_ratio_tolerates_intraday_move(self) -> None:
        """VWAP 与收盘价相差 ±10%（涨跌停）时不得误判单位。"""
        df = pd.DataFrame([
            _row("A", 10.0, 1e6, amount=1e6 * 11.0),   # VWAP 高于收盘 10%
            _row("B", 10.0, 1e6, amount=1e6 * 9.0),    # VWAP 低于收盘 10%
        ])
        assert list(detect_ohlcv_units(df)) == [OHLCV_CONVERTED, OHLCV_CONVERTED]

    def test_degenerate_rows_are_unknown(self) -> None:
        df = pd.DataFrame([
            {"symbol": "A", "close": 0.0, "volume": 100.0, "amount": 1000.0},
            {"symbol": "B", "close": 10.0, "volume": 0.0, "amount": 1000.0},
            {"symbol": "C", "close": 10.0, "volume": 100.0, "amount": float("nan")},
        ])
        assert list(detect_ohlcv_units(df)) == [OHLCV_UNKNOWN] * 3


class TestNormalizeOhlcvUnits:
    def test_raw_rows_are_converted(self) -> None:
        raw = pd.DataFrame([{"symbol": "000001.SZ", "close": 9.80,
                             "volume": 417732.0, "amount": 412636.0}])
        out, diag = normalize_ohlcv_units(raw)
        assert diag["raw"] == 1 and diag["fixed"] == 1
        # 41,773,200 股 × 9.8784 ≈ 4.126 亿元
        assert out.iloc[0]["volume"] == pytest.approx(41_773_200.0)
        assert out.iloc[0]["amount"] == pytest.approx(412_636_000.0)
        assert out.iloc[0]["amount"] / out.iloc[0]["volume"] == pytest.approx(9.878, abs=0.01)

    def test_is_idempotent(self) -> None:
        raw = pd.DataFrame([_raw_row("A", 10.0, 1e6)])
        once, diag1 = normalize_ohlcv_units(raw)
        assert diag1["fixed"] == 1
        twice, diag2 = normalize_ohlcv_units(once)
        pd.testing.assert_frame_equal(once, twice)
        assert diag2["raw"] == 0, "已换算的行不得被二次换算"

    def test_mixed_units_within_one_date(self) -> None:
        """2019 年同日两种单位并存（实测 3596 行里 237 行已换算）。"""
        df = pd.DataFrame([
            _raw_row("RAW", 10.0, 1e6),
            _row("CONV", 10.0, 1e6),
        ])
        out, diag = normalize_ohlcv_units(df)
        assert diag["raw"] == 1 and diag["converted"] == 1
        for i in out.index:
            assert out.at[i, "amount"] / (out.at[i, "volume"] * out.at[i, "close"]) == pytest.approx(1.0)

    def test_does_not_mutate_input(self) -> None:
        raw = pd.DataFrame([_raw_row("A", 10.0, 1e6)])
        before = raw.copy()
        normalize_ohlcv_units(raw)
        pd.testing.assert_frame_equal(raw, before)

    def test_empty_frame(self) -> None:
        out, diag = normalize_ohlcv_units(pd.DataFrame())
        assert out.empty and diag["total"] == 0

    def test_missing_column_is_reported_not_guessed(self) -> None:
        """缺 close 时无法判定单位 —— 必须原样返回并说明，不得瞎猜。"""
        df = pd.DataFrame([{"symbol": "A", "volume": 1e4, "amount": 1e5}])
        out, diag = normalize_ohlcv_units(df)
        pd.testing.assert_frame_equal(out, df)
        assert "error" in diag and diag["fixed"] == 0

    def test_custom_column_names(self) -> None:
        df = pd.DataFrame([_raw_row("A", 10.0, 1e6)].copy())
        df = df.rename(columns={"volume": "vol", "amount": "amt", "close": "px"})
        out, diag = normalize_ohlcv_units(df, volume_col="vol", amount_col="amt", close_col="px")
        assert diag["fixed"] == 1
        assert out.iloc[0]["vol"] == pytest.approx(1e6)
        assert out.iloc[0]["amt"] == pytest.approx(1e7)