easy_tdx_max/tests/unit/test_factor_engine.py

# tests/unit/test_factor_engine.py
"""Test FactorEngine."""
from __future__ import annotations

import numpy as np
import pandas as pd
import pytest

from easy_tdx.factor.base import Factor
from easy_tdx.factor.engine import FactorEngine


def _make_df(n: int = 60, seed: int = 42) -> pd.DataFrame:
    """生成合成 OHLCV 数据。"""
    rng = np.random.default_rng(seed)
    close = 10.0 + np.cumsum(rng.normal(0, 0.5, n))
    close = np.maximum(close, 1.0)
    high = close + rng.uniform(0, 0.5, n)
    low = close - rng.uniform(0, 0.5, n)
    low = np.maximum(low, 0.1)
    open_ = low + rng.uniform(0, high - low, n)
    vol = rng.integers(100_000, 10_000_000, n).astype(float)
    amount = close * vol

    dates = pd.date_range("2024-01-01", periods=n, freq="D")
    return pd.DataFrame({
        "datetime": dates,
        "open": open_,
        "high": high,
        "low": low,
        "close": close,
        "vol": vol,
        "amount": amount,
    })


class _SimpleMomentum(Factor):
    name = "simple_momentum"
    category = "momentum"
    description = "5 日动量"
    inputs = ("close",)

    def compute(self, df: pd.DataFrame) -> pd.Series:
        return df["close"].pct_change(5)


class _SimpleVolatility(Factor):
    name = "simple_volatility"
    category = "volatility"
    description = "5 日波动率"
    inputs = ("close",)

    def compute(self, df: pd.DataFrame) -> pd.Series:
        ret = df["close"].pct_change()
        return ret.rolling(5).std()


class TestComputeSingle:
    def test_single_factor(self):
        engine = FactorEngine()
        df = _make_df()
        result = engine.compute_single(df, [_SimpleMomentum()])
        assert "simple_momentum" in result.columns
        assert len(result) == len(df)

    def test_multiple_factors(self):
        engine = FactorEngine()
        df = _make_df()
        result = engine.compute_single(df, [_SimpleMomentum(), _SimpleVolatility()])
        assert "simple_momentum" in result.columns
        assert "simple_volatility" in result.columns
        assert len(result) == len(df)

    def test_preserves_original_columns(self):
        engine = FactorEngine()
        df = _make_df()
        result = engine.compute_single(df, [_SimpleMomentum()])
        assert "close" in result.columns
        assert "datetime" in result.columns

    def test_unknown_factor_name_raises(self):
        engine = FactorEngine()
        df = _make_df()
        with pytest.raises(ValueError, match="未知因子"):
            engine.compute_single(df, ["nonexistent_factor_xyz"])

    def test_empty_factors_list(self):
        engine = FactorEngine()
        df = _make_df()
        result = engine.compute_single(df, [])
        assert len(result) == len(df)


class TestComputeCrossSection:
    def test_cross_section_basic(self):
        engine = FactorEngine()
        data = {
            "000001": _make_df(60, seed=1),
            "000002": _make_df(60, seed=2),
            "600036": _make_df(60, seed=3),
        }
        result = engine.compute_cross_section(data, [_SimpleMomentum()])
        assert isinstance(result, pd.DataFrame)
        assert "date" in result.columns
        assert "code" in result.columns
        assert "simple_momentum" in result.columns
        assert len(result) == 180  # 60 days × 3 stocks

    def test_cross_section_latest_date(self):
        engine = FactorEngine()
        data = {
            "000001": _make_df(60, seed=1),
            "000002": _make_df(60, seed=2),
        }
        result = engine.compute_cross_section(data, [_SimpleMomentum()], date=None)
        assert len(result) == 2

    def test_cross_section_specific_date(self):
        engine = FactorEngine()
        df = _make_df(60, seed=1)
        data = {"000001": df}
        target_date = int(df["datetime"].iloc[-5].strftime("%Y%m%d"))
        result = engine.compute_cross_section(data, [_SimpleMomentum()], date=target_date)
        assert len(result) == 1
        assert result.iloc[0]["date"] == target_date

    def test_cross_section_empty_data(self):
        engine = FactorEngine()
        result = engine.compute_cross_section({}, [_SimpleMomentum()])
        assert len(result) == 0


class TestComputeForwardReturns:
    def test_forward_returns_basic(self):
        engine = FactorEngine()
        data = {
            "000001": _make_df(60, seed=1),
            "000002": _make_df(60, seed=2),
        }
        result = engine.compute_forward_returns(data, period=5)
        assert "date" in result.columns
        assert "code" in result.columns
        assert "forward_5d" in result.columns
        code_000001 = result[result["code"] == "000001"]
        assert np.isnan(code_000001["forward_5d"].iloc[-1])

    def test_forward_returns_period(self):
        engine = FactorEngine()
        data = {"000001": _make_df(60, seed=1)}
        result = engine.compute_forward_returns(data, period=10)
        assert "forward_10d" in result.columns

    def test_forward_returns_empty(self):
        engine = FactorEngine()
        result = engine.compute_forward_returns({}, period=5)
        assert len(result) == 0


class TestFactorEngineWithBuiltins:
    """FactorEngine 与内置因子的集成测试。"""

    def test_compute_single_with_builtin(self):
        engine = FactorEngine()
        df = _make_df(120)
        result = engine.compute_single(df, ["momentum_20d", "volatility_20d", "rsi_14"])
        assert "momentum_20d" in result.columns
        assert "volatility_20d" in result.columns
        assert "rsi_14" in result.columns
        assert not result["momentum_20d"].iloc[20:25].isna().all()

    def test_cross_section_with_builtins(self):
        engine = FactorEngine()
        data = {
            "000001": _make_df(120, seed=1),
            "000002": _make_df(120, seed=2),
        }
        result = engine.compute_cross_section(data, ["momentum_20d", "sharpe_20d"])
        assert "momentum_20d" in result.columns
        assert "sharpe_20d" in result.columns
        assert len(result) == 240

    def test_forward_returns_with_data(self):
        engine = FactorEngine()
        data = {
            "000001": _make_df(120, seed=1),
        }
        result = engine.compute_forward_returns(data, period=5)
        assert "forward_5d" in result.columns
        assert len(result) == 120
        assert not np.isnan(result["forward_5d"].iloc[50])
        assert np.isnan(result["forward_5d"].iloc[-1])

    def test_all_builtin_factors_compute(self):
        """验证所有内置因子都能无报错地计算。"""
        engine = FactorEngine()
        df = _make_df(200)

        from easy_tdx.factor.builtin import list_factors
        for f_info in list_factors():
            name = f_info["name"]
            result = engine.compute_single(df, [name])
            assert name in result.columns, f"因子 {name} 计算失败"