"""Test built-in factor computation correctness.""" from __future__ import annotations import numpy as np import pandas as pd import pytest from easy_tdx.factor.base import FACTORY_REGISTRY from easy_tdx.factor.builtin import get_factor, list_factors def _make_df(n: int = 120, seed: int = 42) -> pd.DataFrame: """生成合成 OHLCV 数据(120 行,满足所有因子最小窗口)。""" rng = np.random.default_rng(seed) close = 10.0 + np.cumsum(rng.normal(0, 0.3, n)) close = np.maximum(close, 1.0) high = close + rng.uniform(0, 0.3, n) low = close - rng.uniform(0, 0.3, 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, } ) # ── Auto-registration ────────────────────────────────────────────── class TestAutoRegistration: def test_momentum_factors_registered(self): assert "momentum_20d" in FACTORY_REGISTRY assert "momentum_60d" in FACTORY_REGISTRY assert "reversal_5d" in FACTORY_REGISTRY def test_volatility_factors_registered(self): assert "volatility_20d" in FACTORY_REGISTRY assert "atr_14d" in FACTORY_REGISTRY assert "turnover_rate" in FACTORY_REGISTRY def test_quality_factors_registered(self): assert "sharpe_20d" in FACTORY_REGISTRY assert "max_drawdown_20d" in FACTORY_REGISTRY assert "win_rate_20d" in FACTORY_REGISTRY def test_volume_factors_registered(self): assert "obv_trend" in FACTORY_REGISTRY assert "vol_surge" in FACTORY_REGISTRY assert "amount_ma_ratio" in FACTORY_REGISTRY def test_technical_factors_registered(self): assert "macd_hist_signal" in FACTORY_REGISTRY assert "rsi_14" in FACTORY_REGISTRY assert "boll_position" in FACTORY_REGISTRY def test_chanlun_factors_registered(self): assert "chanlun_bi_dir" in FACTORY_REGISTRY assert "chanlun_mmd" in FACTORY_REGISTRY def test_value_factors_registered(self): assert "pe_ratio" in FACTORY_REGISTRY assert "pb_ratio" in FACTORY_REGISTRY def test_total_factor_count(self): assert len(FACTORY_REGISTRY) >= 19 # ── list_factors / get_factor ─────────────────────────────────────── class TestListAndGetFactors: def test_list_factors_returns_all(self): factors = list_factors() assert len(factors) >= 19 for f in factors: assert "name" in f assert "category" in f assert "description" in f def test_get_factor_existing(self): cls = get_factor("momentum_20d") assert cls.name == "momentum_20d" def test_get_factor_nonexistent(self): with pytest.raises(ValueError, match="未知因子"): get_factor("nonexistent") # ── Momentum compute ─────────────────────────────────────────────── class TestMomentumCompute: def test_momentum_20d(self): f = get_factor("momentum_20d")() df = _make_df() result = f.compute(df) assert isinstance(result, pd.Series) assert len(result) == len(df) assert not np.isnan(result.iloc[20]) def test_momentum_60d(self): f = get_factor("momentum_60d")() df = _make_df() result = f.compute(df) assert not np.isnan(result.iloc[60]) def test_reversal_5d_is_negative_return(self): f = get_factor("reversal_5d")() df = _make_df() result = f.compute(df) expected = -df["close"].pct_change(5) pd.testing.assert_series_equal(result, expected, check_names=False) # ── Volatility compute ───────────────────────────────────────────── class TestVolatilityCompute: def test_volatility_20d(self): f = get_factor("volatility_20d")() df = _make_df() result = f.compute(df) assert result.iloc[20] > 0 def test_atr_14d(self): f = get_factor("atr_14d")() df = _make_df() result = f.compute(df) assert result.iloc[14] > 0 def test_turnover_rate(self): f = get_factor("turnover_rate")() df = _make_df() result = f.compute(df) assert result.iloc[40] > 0 # ── Quality compute ──────────────────────────────────────────────── class TestQualityCompute: def test_sharpe_20d(self): f = get_factor("sharpe_20d")() df = _make_df() result = f.compute(df) assert len(result) == len(df) def test_max_drawdown_20d(self): f = get_factor("max_drawdown_20d")() df = _make_df() result = f.compute(df) valid = result.dropna() assert (valid <= 0).all() def test_win_rate_20d(self): f = get_factor("win_rate_20d")() df = _make_df() result = f.compute(df) valid = result.dropna() assert (valid >= 0).all() assert (valid <= 1).all() # ── Volume compute ───────────────────────────────────────────────── class TestVolumeCompute: def test_vol_surge(self): f = get_factor("vol_surge")() df = _make_df() result = f.compute(df) assert result.iloc[20] > 0 def test_amount_ma_ratio(self): f = get_factor("amount_ma_ratio")() df = _make_df() result = f.compute(df) assert len(result) == len(df) # ── Technical compute ────────────────────────────────────────────── class TestTechnicalCompute: def test_rsi_14_range(self): f = get_factor("rsi_14")() df = _make_df() result = f.compute(df) valid = result.dropna() assert (valid >= -1).all() assert (valid <= 1).all() def test_boll_position_range(self): f = get_factor("boll_position")() df = _make_df() result = f.compute(df) valid = result.dropna() assert (valid >= 0).all() assert (valid <= 1).all()