ProStock/tests/test_financial_price_merge.py

"""财务数据与行情数据拼接测试。

测试场景：
1. 普通财务数据：正常公告，之后无修改
2. 隔日修改：公告后几天发布修正版
3. 当日修改：同一天发布多版，取 update_flag=1 的
4. 边界条件：财务数据缺失、行情数据早于最早财务数据
"""

import polars as pl
from datetime import date
from src.data.financial_loader import FinancialLoader


def create_mock_price_data() -> pl.DataFrame:
    """创建模拟行情数据。"""
    return pl.DataFrame(
        {
            "ts_code": ["000001.SZ"] * 12,
            "trade_date": [
                "20240101",
                "20240102",
                "20240103",
                "20240104",
                "20240105",
                "20240108",
                "20240109",
                "20240110",
                "20240111",
                "20240112",
                # 添加2024-04-30之后的日期，用于测试同日不同报告期场景
                "20240501",
                "20240502",
            ],
            "close": [
                10.0,
                10.2,
                10.3,
                10.1,
                10.5,
                10.6,
                10.4,
                10.7,
                10.8,
                10.9,
                11.0,
                11.1,
            ],
        }
    )


def create_mock_financial_data() -> pl.DataFrame:
    """创建模拟财务数据（覆盖多种场景）。

    场景说明：
    1. 2024-01-02 发布 2023Q3 报告（end_date=20230930）
    2. 2024-01-02 发布 2023Q3 更正版（update_flag=1）
    3. 2024-04-30 同时发布 2023年报（end_date=20231231）和 2024Q1季报（end_date=20240331）
    4. 2024-04-30 发布 2023年报更正版

    预期结果：
    - 2024-01-02 保留 2023Q3 更正版
    - 2024-04-30 保留 2024Q1 季报（end_date 最新）

    注意：f_ann_date 必须是 Date 类型（与数据库保持一致）。
    """
    return pl.DataFrame(
        {
            "ts_code": [
                "000001.SZ",
                "000001.SZ",
                "000001.SZ",
                "000001.SZ",
                "000001.SZ",
            ],
            "f_ann_date": [
                date(2024, 1, 2),
                date(2024, 1, 2),  # 同日多版
                date(2024, 4, 30),
                date(2024, 4, 30),
                date(2024, 4, 30),  # 同日不同报告期
            ],
            "end_date": [
                "20230930",
                "20230930",  # 2023Q3
                "20231231",
                "20240331",
                "20231231",  # 年报和季报同一天发布
            ],
            "report_type": [1, 1, 1, 1, 1],  # 整数类型（与数据库一致）
            "update_flag": [0, 1, 0, 0, 1],  # 年报也有更正版
            "net_profit": [
                1000000.0,
                1100000.0,  # 2023Q3
                5000000.0,
                1500000.0,
                5500000.0,  # 年报更正后550万，季报150万
            ],
            "revenue": [
                5000000.0,
                5200000.0,  # 2023Q3
                20000000.0,
                8000000.0,
                22000000.0,
            ],
        }
    )


def test_financial_data_cleaning():
    """测试财务数据清洗逻辑 - 确保同日多报告期时选 end_date 最新的。"""
    print("=== 测试 1: 财务数据清洗 ===")

    df_finance = create_mock_financial_data()
    print("原始财务数据：")
    print(df_finance)

    loader = FinancialLoader()

    # 手动执行新的清洗逻辑
    df = df_finance.filter(pl.col("report_type") == 1)

    # 添加辅助列
    df = df.with_columns(
        [
            pl.col("end_date").cast(pl.Int32).alias("end_date_int"),
            pl.col("update_flag")
            .fill_null("0")
            .cast(pl.Int32, strict=False)
            .fill_null(0)
            .alias("update_flag_int"),
        ]
    )

    # 确定性排序
    df = df.sort(["ts_code", "f_ann_date", "end_date_int", "update_flag_int"])

    # 累积最大报告期
    df = df.with_columns(
        pl.col("end_date_int").cum_max().over("ts_code").alias("max_end_date_seen")
    )

    # 过滤历史包袱
    df = df.filter(pl.col("end_date_int") == pl.col("max_end_date_seen"))

    # 去重保留最后一条（end_date 最大的）
    df = df.unique(subset=["ts_code", "f_ann_date"], keep="last")

    # 清理辅助列
    df = df.drop(["end_date_int", "update_flag_int", "max_end_date_seen"])
    df = df.sort(["ts_code", "f_ann_date"])

    print("\n清洗后的财务数据：")
    print(df)

    # 验证：应该有2条记录（2024-01-02 和 2024-04-30）
    assert len(df) == 2, f"清洗后应该有2条记录，实际有 {len(df)} 条"

    # 验证：2024-01-02 的 end_date 应该是 20230930
    row_jan02 = df.filter(pl.col("f_ann_date") == date(2024, 1, 2))
    assert len(row_jan02) == 1
    assert row_jan02["end_date"][0] == "20230930"
    assert row_jan02["update_flag"][0] == 1
    print("[验证 1] 2024-01-02 正确保留了 2023Q3 更正版")

    # 验证：2024-04-30 应该保留 2024Q1（end_date=20240331），而不是年报
    row_apr30 = df.filter(pl.col("f_ann_date") == date(2024, 4, 30))
    assert len(row_apr30) == 1
    assert row_apr30["end_date"][0] == "20240331", (
        f"2024-04-30 应该保留 end_date 最新的 20240331，"
        f"实际为 {row_apr30['end_date'][0]}"
    )
    assert row_apr30["net_profit"][0] == 1500000.0
    print("[验证 2] 2024-04-30 正确保留了 2024Q1 季报（end_date 最新）")

    print("\n[通过] 财务数据清洗测试通过！")
    return df


def test_financial_price_merge():
    """测试财务数据拼接逻辑（无未来函数验证）。"""
    print("\n=== 测试 2: 财务数据与行情数据拼接 ===")

    df_price = create_mock_price_data()
    df_finance_raw = create_mock_financial_data()

    loader = FinancialLoader()

    # 步骤1: 清洗财务数据（手动执行新的清洗逻辑）
    # 注意：f_ann_date 已经是 Date 类型，不需要转换
    df_finance = df_finance_raw.filter(pl.col("report_type") == 1)

    # 添加辅助列
    df_finance = df_finance.with_columns(
        [
            pl.col("end_date").cast(pl.Int32).alias("end_date_int"),
            pl.col("update_flag")
            .fill_null("0")
            .cast(pl.Int32, strict=False)
            .fill_null(0)
            .alias("update_flag_int"),
        ]
    )

    # 确定性排序
    df_finance = df_finance.sort(
        ["ts_code", "f_ann_date", "end_date_int", "update_flag_int"]
    )

    # 累积最大报告期
    df_finance = df_finance.with_columns(
        pl.col("end_date_int").cum_max().over("ts_code").alias("max_end_date_seen")
    )

    # 过滤历史包袱
    df_finance = df_finance.filter(
        pl.col("end_date_int") == pl.col("max_end_date_seen")
    )

    # 去重保留最后一条（end_date 最大的）
    df_finance = df_finance.unique(subset=["ts_code", "f_ann_date"], keep="last")

    # 清理辅助列
    df_finance = df_finance.drop(
        ["end_date_int", "update_flag_int", "max_end_date_seen"]
    )
    df_finance = df_finance.sort(["ts_code", "f_ann_date"])

    print("清洗后的财务数据：")
    print(df_finance)

    # 步骤2: 转换行情数据日期为 Date 类型
    df_price = df_price.with_columns(
        [pl.col("trade_date").str.strptime(pl.Date, "%Y%m%d").alias("trade_date")]
    )
    df_price = df_price.sort(["ts_code", "trade_date"])

    # 步骤3: 拼接
    financial_cols = ["net_profit", "revenue"]
    merged = loader.merge_financial_with_price(df_price, df_finance, financial_cols)

    # 步骤4: 转回字符串格式
    merged = merged.with_columns(
        [pl.col("trade_date").dt.strftime("%Y%m%d").alias("trade_date")]
    )

    print("\n拼接结果：")
    print(merged)

    # 验证无未来函数：
    # 20240101 之前不应有 2023Q3 数据（因为 20240102 才公告）
    jan01 = merged.filter(pl.col("trade_date") == "20240101")
    assert jan01["net_profit"].is_null().all(), (
        "2024-01-01 不应有 2023Q3 数据（尚未公告）"
    )
    print("[验证 1] 2024-01-01 net_profit 为 null - 正确（公告前无数据）")

    # 20240102 及之后应该看到 net_profit=1100000（update_flag=1 的版本）
    jan02 = merged.filter(pl.col("trade_date") == "20240102")
    assert jan02["net_profit"][0] == 1100000.0, "2024-01-02 应使用 update_flag=1 的数据"
    print("[验证 2] 2024-01-02 net_profit=1100000 - 正确（使用 update_flag=1）")

    # 20240104 应延续使用 2023Q3 数据
    jan04 = merged.filter(pl.col("trade_date") == "20240104")
    assert jan04["net_profit"][0] == 1100000.0, "2024-01-04 应延续使用 2023Q3 数据"
    print("[验证 3] 2024-01-04 net_profit=1100000 - 正确（延续使用）")

    # 20240110 应延续使用 2023Q3 数据（2024-04-30 还未公告）
    jan10 = merged.filter(pl.col("trade_date") == "20240110")
    assert jan10["net_profit"][0] == 1100000.0, "2024-01-10 应延续使用 2023Q3 数据"
    print("[验证 4] 2024-01-10 net_profit=1100000 - 正确（延续使用 2023Q3）")

    # 20240112 应继续延续使用 2023Q3 数据
    jan12 = merged.filter(pl.col("trade_date") == "20240112")
    assert jan12["net_profit"][0] == 1100000.0, "2024-01-12 应继续使用 2023Q3 数据"
    print("[验证 5] 2024-01-12 net_profit=1100000 - 正确（延续使用 2023Q3）")

    # 20240501 应切换到 2024Q1 数据（2024-04-30 已公告，且选择 end_date 最新的）
    may01 = merged.filter(pl.col("trade_date") == "20240501")
    assert may01["net_profit"][0] == 1500000.0, "2024-05-01 应切换到 2024Q1 数据"
    print(
        "[验证 6] 2024-05-01 net_profit=1500000 - 正确（切换到 2024Q1，end_date 最新）"
    )

    print("\n[通过] 所有验证通过，无未来函数！")
    return merged


def test_empty_financial_data():
    """测试财务数据为空的情况。"""
    print("\n=== 测试 3: 空财务数据场景 ===")

    df_price = create_mock_price_data()
    df_empty = pl.DataFrame()

    loader = FinancialLoader()

    # 转换行情数据日期为 Date 类型
    df_price = df_price.with_columns(
        [pl.col("trade_date").str.strptime(pl.Date, "%Y%m%d").alias("trade_date")]
    )
    df_price = df_price.sort(["ts_code", "trade_date"])

    # 拼接空财务数据
    merged = loader.merge_financial_with_price(df_price, df_empty, ["net_profit"])

    # 转回字符串格式
    merged = merged.with_columns(
        [pl.col("trade_date").dt.strftime("%Y%m%d").alias("trade_date")]
    )

    # 验证财务列为空
    assert merged["net_profit"].is_null().all(), (
        "财务数据为空时，net_profit 应全为 null"
    )

    print("空财务数据拼接结果：")
    print(merged)
    print("\n[通过] 空财务数据场景测试通过！")


def run_all_tests():
    """运行所有测试。"""
    print("开始运行财务数据拼接功能测试...\n")
    print("=" * 60)

    try:
        # 测试 1: 数据清洗
        test_financial_data_cleaning()

        # 测试 2: 数据拼接
        test_financial_price_merge()

        # 测试 3: 空数据场景
        test_empty_financial_data()

        print("\n" + "=" * 60)
        print("所有测试通过！")
        print("=" * 60)

    except AssertionError as e:
        print(f"\n[失败] 测试断言失败: {e}")
        raise
    except Exception as e:
        print(f"\n[错误] 测试执行出错: {e}")
        raise


if __name__ == "__main__":
    run_all_tests()