easy-tdx/tests/unit/test_optimizer.py

"""单元测试：参数网格寻优器."""

from __future__ import annotations

import numpy as np
import pandas as pd
import pytest

from easy_tdx.backtest.optimizer import (
    GridPointResult,
    OptimizeResult,
    ParamGridOptimizer,
)


def _make_df(n: int = 150, seed: int = 42) -> pd.DataFrame:
    """生成带趋势的合成 OHLCV（确保均线策略能产生交易）。"""
    rng = np.random.default_rng(seed)
    close = 10.0 + np.cumsum(rng.normal(0, 0.3, n) + 0.05)
    return pd.DataFrame(
        {
            "datetime": pd.date_range("2024-01-01", periods=n, freq="B"),
            "open": close - 0.1,
            "high": close + 0.2,
            "low": close - 0.2,
            "close": close,
            "vol": rng.integers(1000, 10000, n).astype(float),
            "amount": close * 5000,
        }
    )


class TestParamGridOptimizer:
    """寻优器核心逻辑."""

    def test_grid_enumeration(self) -> None:
        """网格点数应等于笛卡尔积大小。"""
        opt = ParamGridOptimizer(
            strategy_name="ma_cross",
            param_grid={"fast": [5, 10], "slow": [15, 20, 30]},
            df=_make_df(),
        )
        result = opt.run()
        assert len(result.results) == 6  # 2 × 3

    def test_results_sorted_by_return_descending(self) -> None:
        """结果应按 total_return 降序排列。"""
        opt = ParamGridOptimizer(
            strategy_name="ma_cross",
            param_grid={"fast": [5, 10, 20], "slow": [15, 20, 30]},
            df=_make_df(),
        )
        result = opt.run()
        returns = [r.total_return for r in result.results]
        assert returns == sorted(returns, reverse=True)

    def test_best_is_first_result(self) -> None:
        """best 应是 results[0]。"""
        opt = ParamGridOptimizer(
            strategy_name="ma_cross",
            param_grid={"fast": [5, 10], "slow": [20, 30]},
            df=_make_df(),
        )
        result = opt.run()
        assert result.best is not None
        assert result.best.params == result.results[0].params
        assert result.best.total_return == result.results[0].total_return

    def test_heatmap_2_params(self) -> None:
        """2 参数应生成热力图矩阵。"""
        opt = ParamGridOptimizer(
            strategy_name="ma_cross",
            param_grid={"fast": [5, 10, 20], "slow": [15, 20, 30]},
            df=_make_df(),
        )
        result = opt.run()
        assert result.heatmap is not None
        assert result.heatmap["x_name"] == "fast"
        assert result.heatmap["y_name"] == "slow"
        assert len(result.heatmap["x"]) == 3
        assert len(result.heatmap["y"]) == 3
        assert len(result.heatmap["data"]) == 9  # 3×3

    def test_no_heatmap_for_single_param(self) -> None:
        """1 参数时 heatmap 应为 None。"""
        opt = ParamGridOptimizer(
            strategy_name="rsi_reversal",
            param_grid={"n": [7, 14, 21]},
            df=_make_df(),
        )
        result = opt.run()
        assert result.heatmap is None
        assert len(result.results) == 3

    def test_grid_size_limit_exceeded(self) -> None:
        """超过 MAX_GRID_POINTS 应抛 ValueError。"""
        big_grid = {f"p{i}": list(range(5)) for i in range(6)}  # 5^6 = 15625
        with pytest.raises(ValueError, match="超过上限"):
            ParamGridOptimizer(
                strategy_name="ma_cross",
                param_grid=big_grid,
                df=_make_df(),
            )

    def test_empty_value_list_rejected(self) -> None:
        """空取值列表应抛 ValueError。"""
        with pytest.raises(ValueError, match="空取值列表"):
            ParamGridOptimizer(
                strategy_name="ma_cross",
                param_grid={"fast": [], "slow": [20]},
                df=_make_df(),
            )

    def test_to_dict_serializable(self) -> None:
        """to_dict 应返回 JSON 兼容结构。"""
        opt = ParamGridOptimizer(
            strategy_name="ma_cross",
            param_grid={"fast": [5, 10], "slow": [20, 30]},
            df=_make_df(),
        )
        result = opt.run()
        d = result.to_dict()

        assert d["strategy"] == "ma_cross"
        assert d["param_names"] == ["fast", "slow"]
        assert len(d["results"]) == 4
        assert d["best"] is not None
        assert "params" in d["best"]
        assert "total_return" in d["best"]

    def test_to_dict_cleans_nan(self) -> None:
        """NaN 指标应被清洗为 None（JSON 兼容）。"""
        # 构造含 NaN 的结果（无交易的参数组合 sharpe 可能 NaN）
        result = OptimizeResult(
            strategy="test",
            param_names=["n"],
            results=[
                GridPointResult(params={"n": 1}, total_return=float("nan"), sharpe=float("inf")),
            ],
        )
        d = result.to_dict()
        assert d["results"][0]["total_return"] is None
        assert d["results"][0]["sharpe"] is None

    def test_unknown_strategy_raises(self) -> None:
        """未知策略应在 run() 时抛 KeyError。"""
        opt = ParamGridOptimizer(
            strategy_name="nope",
            param_grid={"x": [1]},
            df=_make_df(),
        )
        with pytest.raises(KeyError):
            opt.run()