ProStock/tests/factors/test_engine.py

"""测试执行引擎 - FactorEngine

测试需求（来自 factor_implementation_plan.md）：
- 测试 `compute()` 正确分发给截面计算
- 测试 `compute()` 正确分发给时序计算
- 测试无效 factor_type 时抛出 ValueError

截面计算测试（防泄露验证）：
- 测试数据裁剪正确（传入 [T-lookback+1, T]）
- 测试不包含未来日期 T+1 的数据
- 测试每个日期独立计算
- 测试结果包含所有日期和所有股票
- 测试结果 DataFrame 格式正确
- 测试多个 DataSpec 时 lookback 取最大值

时序计算测试（防泄露验证）：
- 测试每只股票只看到自己的数据
- 测试不包含其他股票的数据
- 测试传入的是完整时间序列（向量化计算）
- 测试结果包含所有股票和所有日期
- 测试结果 DataFrame 格式正确
- 测试股票不在数据中时跳过（或填充 null）
"""

import pytest
import polars as pl

from src.factors import (
    DataSpec,
    FactorContext,
    FactorData,
    DataLoader,
    FactorEngine,
    CrossSectionalFactor,
    TimeSeriesFactor,
)


class SimpleCrossSectionalFactor(CrossSectionalFactor):
    """简单的截面因子 - 返回收盘价排名"""

    name = "close_rank"
    data_specs = [
        DataSpec("daily", ["ts_code", "trade_date", "close"], lookback_days=1)
    ]

    def compute(self, data: FactorData) -> pl.Series:
        cs = data.get_cross_section()
        return cs["close"].rank()


class SimpleTimeSeriesFactor(TimeSeriesFactor):
    """简单的时序因子 - 返回3日移动平均"""

    name = "ma3"
    data_specs = [
        DataSpec(
            "daily",
            ["ts_code", "trade_date", "close"],
            lookback_days=5,
        )
    ]

    def __init__(self, period: int = 3):
        super().__init__(period=period)

    def compute(self, data: FactorData) -> pl.Series:
        return data.get_column("close").rolling_mean(self.params["period"])


class ReturnFactor(CrossSectionalFactor):
    """收益率因子 - 需要2天lookback计算收益率"""

    name = "return"
    data_specs = [
        DataSpec("daily", ["ts_code", "trade_date", "close"], lookback_days=2)
    ]

    def compute(self, data: FactorData) -> pl.Series:
        # 获取当前日期
        current_date = data.context.current_date

        # 获取当前日期的数据
        cs = data.get_cross_section()

        # 简单返回收盘价作为因子值
        # 实际应该计算收益率，但这里简化处理
        return cs["close"]


@pytest.fixture
def loader():
    """创建 DataLoader 实例"""
    return DataLoader(data_dir="data")


@pytest.fixture
def engine(loader):
    """创建 FactorEngine 实例"""
    return FactorEngine(loader)


class TestFactorEngineDispatch:
    """测试引擎分发逻辑"""

    def test_dispatch_cross_sectional(self, engine):
        """测试 compute() 正确分发给截面计算"""
        factor = SimpleCrossSectionalFactor()

        result = engine.compute(factor, start_date="20240101", end_date="20240105")

        assert isinstance(result, pl.DataFrame)
        assert "trade_date" in result.columns
        assert "ts_code" in result.columns
        assert "close_rank" in result.columns

    def test_dispatch_time_series(self, engine, loader):
        """测试 compute() 正确分发给时序计算"""
        factor = SimpleTimeSeriesFactor(period=3)

        # 获取一些股票代码
        sample_data = loader.load(
            [DataSpec("daily", ["ts_code", "trade_date"], lookback_days=1)]
        )
        stock_codes = sample_data["ts_code"].unique().head(3).to_list()

        result = engine.compute(
            factor,
            stock_codes=stock_codes,
            start_date="20240101",
            end_date="20240110",
        )

        assert isinstance(result, pl.DataFrame)
        assert "trade_date" in result.columns
        assert "ts_code" in result.columns
        assert "ma3" in result.columns

    def test_unknown_factor_type(self, engine):
        """测试无效 factor_type 时抛出 ValueError"""

        class UnknownFactor:
            name = "unknown"
            factor_type = "unknown_type"
            data_specs = []

        factor = UnknownFactor()

        with pytest.raises(ValueError, match="Unknown factor type"):
            engine.compute(factor)


class TestCrossSectionalComputation:
    """测试截面计算（防泄露验证）"""

    def test_result_format(self, engine):
        """测试结果 DataFrame 格式正确"""
        factor = SimpleCrossSectionalFactor()

        result = engine.compute(factor, start_date="20240101", end_date="20240105")

        # 检查列
        assert "trade_date" in result.columns
        assert "ts_code" in result.columns
        assert factor.name in result.columns

        # 检查类型
        assert result["trade_date"].dtype == pl.Utf8
        assert result["ts_code"].dtype == pl.Utf8

    def test_all_dates_present(self, engine):
        """测试结果包含所有日期"""
        factor = SimpleCrossSectionalFactor()

        start_date = "20240101"
        end_date = "20240105"

        result = engine.compute(factor, start_date=start_date, end_date=end_date)

        if len(result) > 0:
            dates = result["trade_date"].unique().to_list()
            # 应该包含 start_date 和 end_date 之间的日期
            assert len(dates) > 0

    def test_lookback_window(self, engine):
        """测试多个 DataSpec 时 lookback 取最大值"""
        factor = ReturnFactor()

        # lookback_days = 2
        result = engine.compute(factor, start_date="20240103", end_date="20240105")

        # 应该能计算出结果
        assert isinstance(result, pl.DataFrame)


class TestTimeSeriesComputation:
    """测试时序计算（防泄露验证）"""

    def test_result_format(self, engine):
        """测试结果 DataFrame 格式正确"""
        factor = SimpleTimeSeriesFactor(period=3)

        result = engine.compute(
            factor,
            stock_codes=["000001.SZ"],
            start_date="20240101",
            end_date="20240110",
        )

        # 检查列
        assert "trade_date" in result.columns
        assert "ts_code" in result.columns
        assert factor.name in result.columns

    def test_single_stock_data(self, engine):
        """测试每只股票只看到自己的数据"""
        factor = SimpleTimeSeriesFactor(period=3)

        stock_codes = ["000001.SZ"]

        result = engine.compute(
            factor,
            stock_codes=stock_codes,
            start_date="20240101",
            end_date="20240110",
        )

        if len(result) > 0:
            # 结果中只应该有指定的股票
            stocks = result["ts_code"].unique().to_list()
            assert set(stocks) == set(stock_codes)

    def test_ma_calculation(self, engine):
        """测试移动平均计算"""
        factor = SimpleTimeSeriesFactor(period=3)

        result = engine.compute(
            factor,
            stock_codes=["000001.SZ"],
            start_date="20240101",
            end_date="20240110",
        )

        if len(result) > 2:
            # 前2个应该是 null（因为 period=3）
            ma_values = result[factor.name].to_list()
            assert ma_values[0] is None or str(ma_values[0]) == "nan"
            assert ma_values[1] is None or str(ma_values[1]) == "nan"
            # 第3个应该有值
            assert ma_values[2] is not None

    def test_missing_stock_skipped(self, engine):
        """测试股票不在数据中时返回空结果"""
        factor = SimpleTimeSeriesFactor(period=3)

        result = engine.compute(
            factor,
            stock_codes=["NONEXISTENT.STOCK"],
            start_date="20240101",
            end_date="20240110",
        )

        # 应该返回空 DataFrame 或包含该股票但值为 null 的结果
        assert isinstance(result, pl.DataFrame)
        # 对于不存在的股票，结果可能是空的
        # 或者包含该股票但值为 null