ProStock/tests/factors/test_data_loader.py

"""测试数据加载器 - DataLoader

测试需求（来自 factor_implementation_plan.md）：
- 测试从 DuckDB 加载数据
- 测试从多个查询加载并合并
- 测试列选择（只加载需要的列）
- 测试缓存机制（第二次加载更快）
- 测试 clear_cache() 清空缓存
- 测试按 date_range 过滤
- 测试表不存在时的处理
- 测试列不存在时抛出 KeyError

使用 3 个月的真实数据进行测试 (2024年1月-3月)
"""

import pytest
import polars as pl
import pandas as pd
from pathlib import Path

from src.factors import DataSpec, DataLoader


class TestDataLoaderBasic:
    """测试 DataLoader 基本功能"""

    # 测试数据时间范围：3个月
    TEST_START_DATE = "20240101"
    TEST_END_DATE = "20240331"

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

    def test_init(self):
        """测试初始化"""
        loader = DataLoader(data_dir="data")
        assert loader.data_dir == Path("data")
        assert loader._cache == {}

    def test_load_single_source(self, loader):
        """测试从 DuckDB 加载数据"""
        specs = [
            DataSpec(
                source="daily",
                columns=["ts_code", "trade_date", "close"],
                lookback_days=1,
            )
        ]

        # 使用 3 个月日期范围限制数据量
        df = loader.load(specs, date_range=(self.TEST_START_DATE, self.TEST_END_DATE))

        assert isinstance(df, pl.DataFrame)
        assert len(df) > 0
        assert "ts_code" in df.columns
        assert "trade_date" in df.columns
        assert "close" in df.columns

    def test_load_with_date_range(self, loader):
        """测试加载特定日期范围（3个月）"""
        specs = [
            DataSpec(
                source="daily",
                columns=["ts_code", "trade_date", "close", "open", "high", "low"],
                lookback_days=1,
            )
        ]

        df = loader.load(specs, date_range=(self.TEST_START_DATE, self.TEST_END_DATE))

        assert isinstance(df, pl.DataFrame)
        assert len(df) > 0

        # 验证日期范围
        if len(df) > 0:
            dates = df["trade_date"].to_list()
            assert all(self.TEST_START_DATE <= d <= self.TEST_END_DATE for d in dates)
            print(f"[TEST] Loaded {len(df)} rows from {min(dates)} to {max(dates)}")

    def test_load_multiple_specs(self, loader):
        """测试从多个 DataSpec 加载并合并"""
        specs = [
            DataSpec(
                source="daily",
                columns=["ts_code", "trade_date", "close"],
                lookback_days=1,
            ),
            DataSpec(
                source="daily",
                columns=["ts_code", "trade_date", "open", "high", "low"],
                lookback_days=1,
            ),
        ]

        df = loader.load(specs, date_range=(self.TEST_START_DATE, self.TEST_END_DATE))

        assert isinstance(df, pl.DataFrame)
        assert len(df) > 0
        # 应该包含所有列
        assert set(df.columns) >= {
            "ts_code",
            "trade_date",
            "close",
            "open",
            "high",
            "low",
        }

    def test_column_selection(self, loader):
        """测试列选择（只加载需要的列）"""
        specs = [
            DataSpec(
                source="daily",
                columns=["ts_code", "trade_date", "close"],
                lookback_days=1,
            )
        ]

        df = loader.load(specs, date_range=(self.TEST_START_DATE, self.TEST_END_DATE))

        # 只应该有 3 列
        assert set(df.columns) == {"ts_code", "trade_date", "close"}

    def test_date_range_filter(self, loader):
        """测试按 date_range 过滤 - 使用3个月数据的不同子集"""
        specs = [
            DataSpec(
                source="daily",
                columns=["ts_code", "trade_date", "close"],
                lookback_days=1,
            )
        ]

        # 加载完整的3个月数据
        df_all = loader.load(
            specs, date_range=(self.TEST_START_DATE, self.TEST_END_DATE)
        )
        total_rows = len(df_all)

        # 清空缓存，重新加载1个月数据
        loader.clear_cache()
        df_filtered = loader.load(specs, date_range=("20240101", "20240131"))

        # 过滤后的数据应该更少或相等
        assert len(df_filtered) <= total_rows

        # 所有日期都应该在范围内
        if len(df_filtered) > 0:
            dates = df_filtered["trade_date"].to_list()
            assert all("20240101" <= d <= "20240131" for d in dates)


class TestDataLoaderCache:
    """测试 DataLoader 缓存机制"""

    TEST_START_DATE = "20240101"
    TEST_END_DATE = "20240331"

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

    def test_cache_populated(self, loader):
        """测试加载后缓存被填充"""
        specs = [
            DataSpec(
                source="daily",
                columns=["ts_code", "trade_date", "close"],
                lookback_days=1,
            )
        ]

        # 第一次加载
        loader.load(specs, date_range=(self.TEST_START_DATE, self.TEST_END_DATE))

        # 检查缓存
        assert len(loader._cache) > 0

    def test_cache_used(self, loader):
        """测试第二次加载使用缓存（更快）"""
        import time

        specs = [
            DataSpec(
                source="daily",
                columns=["ts_code", "trade_date", "close"],
                lookback_days=1,
            )
        ]

        # 第一次加载
        start = time.time()
        df1 = loader.load(specs, date_range=(self.TEST_START_DATE, self.TEST_END_DATE))
        time1 = time.time() - start

        # 第二次加载（应该使用缓存）
        start = time.time()
        df2 = loader.load(specs, date_range=(self.TEST_START_DATE, self.TEST_END_DATE))
        time2 = time.time() - start

        # 数据应该相同
        assert df1.shape == df2.shape

        # 第二次应该更快
        print(f"[TEST] First load: {time1:.3f}s, cached load: {time2:.3f}s")
        assert time2 < time1, "Cached load should be faster"

    def test_clear_cache(self, loader):
        """测试 clear_cache() 清空缓存"""
        specs = [
            DataSpec(
                source="daily",
                columns=["ts_code", "trade_date", "close"],
                lookback_days=1,
            )
        ]

        # 加载数据
        loader.load(specs, date_range=(self.TEST_START_DATE, self.TEST_END_DATE))
        assert len(loader._cache) > 0

        # 清空缓存
        loader.clear_cache()
        assert len(loader._cache) == 0

    def test_cache_info(self, loader):
        """测试 get_cache_info()"""
        specs = [
            DataSpec(
                source="daily",
                columns=["ts_code", "trade_date", "close"],
                lookback_days=1,
            )
        ]

        # 加载前
        info_before = loader.get_cache_info()
        assert info_before["entries"] == 0

        # 加载后
        loader.load(specs, date_range=(self.TEST_START_DATE, self.TEST_END_DATE))
        info_after = loader.get_cache_info()
        assert info_after["entries"] > 0
        assert info_after["total_rows"] > 0


class TestDataLoaderErrors:
    """测试 DataLoader 错误处理"""

    def test_table_not_exists(self):
        """测试表不存在时的处理"""
        loader = DataLoader(data_dir="data")
        specs = [
            DataSpec(
                source="nonexistent_table",
                columns=["ts_code", "trade_date", "close"],
                lookback_days=1,
            )
        ]

        # 应该返回空 DataFrame 或抛出异常
        with pytest.raises(Exception):
            loader.load(specs)

    def test_column_not_found(self):
        """测试列不存在时抛出 KeyError"""
        loader = DataLoader(data_dir="data")
        specs = [
            DataSpec(
                source="daily",
                columns=["ts_code", "trade_date", "nonexistent_column"],
                lookback_days=1,
            )
        ]

        with pytest.raises(KeyError, match="nonexistent_column"):
            loader.load(specs)


if __name__ == "__main__":
    pytest.main([__file__, "-v", "-s"])