"""tests/portfolio/test_migrate_announce_union.py TDD for migrate_constituent._read_announce_union_aggregated() 聚合逻辑(全集型): - ever_codes = announce_union 所有 distinct code (add/remove/initial/current 任一) - current_codes = announce_union 中 adjust_type='current' 的 code (∪ snapshot 兜底) - in_current = code in current_codes - was_removed = not in_current (曾经入选但当前不在) 对应 plan: docs/superpowers/plans/2026-07-23-csi1000-constituent-backfill.md Task 2 """ import os import sys from pathlib import Path import pandas as pd import pytest # 让测试能 import scripts/... ROOT = Path(__file__).resolve().parents[2] sys.path.insert(0, str(ROOT)) sys.path.insert(0, str(ROOT / "scripts" / "data_platform")) # ---------- import target ---------- @pytest.fixture(scope="module") def migrate_module(): """import migrate_constituent 模块(不走 main, 只用函数)""" import importlib.util spec = importlib.util.spec_from_file_location( "migrate_constituent_for_test", ROOT / "scripts" / "data_platform" / "migrate_constituent.py", ) mod = importlib.util.module_from_spec(spec) spec.loader.exec_module(mod) return mod # ---------- 合成 announce_union ---------- def _make_announce_df(rows): """rows: list of (code, code_name, adjust_type)""" return pd.DataFrame( [{"updateDate": "2024-01-01", "index_code": "000852", "code": c, "code_name": n, "adjust_type": t, "notice_id": 1, "source": "synthetic.xlsx"} for c, n, t in rows] ) # ============================================================ # Unit: 聚合逻辑(ever/in_current/was_removed) # ============================================================ def test_aggregation_basic(migrate_module, tmp_path): """announce(add A,B + remove C) + current(A,B,D) -> ever={A,B,C,D}, in_current={A,B,D}, was_removed={C} """ # build announce_union parquet (含 current 行) ann = _make_announce_df([ ("000001", "A", "add"), ("000002", "B", "add"), ("000003", "C", "remove"), ("000001", "A", "current"), ("000002", "B", "current"), ("000004", "D", "current"), ]) hist = tmp_path / "hist" hist.mkdir() ann.to_parquet(hist / "000852_announce_union.parquet", index=False) df = migrate_module._read_announce_union_aggregated(str(hist), ["000852"]) assert len(df) == 4 # ever set 4 stocks # index_code 全是 000852 assert set(df["index_code"]) == {"000852"} # ever = {A,B,C,D} assert set(df["code"]) == {"000001", "000002", "000003", "000004"} # in_current = {A,B,D}, was_removed = {C} in_cur = set(df[df["in_current"] == 1]["code"]) removed = set(df[df["was_removed"] == 1]["code"]) assert in_cur == {"000001", "000002", "000004"} assert removed == {"000003"} # 不存在既 in_current 又 was_removed 的行 assert ((df["in_current"] == 1) & (df["was_removed"] == 1)).sum() == 0 # 每行至少一个标记(全集型的硬条件) assert ((df["in_current"] == 1) | (df["was_removed"] == 1)).sum() == 4 def test_aggregation_no_current_rows_uses_snapshot(migrate_module, tmp_path): """announce_union 只有 add/remove 时, 从 snapshot.parquet 兜底 current_codes""" ann = _make_announce_df([ ("000001", "A", "add"), ("000002", "B", "add"), ("000003", "C", "remove"), ]) snap = pd.DataFrame({ "index_code": ["000852", "000852"], "code": ["000001", "000002"], "code_name": ["A", "B"], }) hist = tmp_path / "hist" hist.mkdir() ann.to_parquet(hist / "000852_announce_union.parquet", index=False) snap.to_parquet(hist / "000852_snapshot.parquet", index=False) df = migrate_module._read_announce_union_aggregated(str(hist), ["000852"]) assert set(df["code"]) == {"000001", "000002", "000003"} in_cur = set(df[df["in_current"] == 1]["code"]) assert in_cur == {"000001", "000002"} # snapshot 兜底 removed = set(df[df["was_removed"] == 1]["code"]) assert removed == {"000003"} def test_aggregation_idempotent(migrate_module, tmp_path): """跑两次结果一致(幂等)""" ann = _make_announce_df([ ("000001", "A", "add"), ("000002", "B", "remove"), ("000001", "A", "current"), ]) hist = tmp_path / "hist" hist.mkdir() ann.to_parquet(hist / "000852_announce_union.parquet", index=False) df1 = migrate_module._read_announce_union_aggregated(str(hist), ["000852"]) df2 = migrate_module._read_announce_union_aggregated(str(hist), ["000852"]) pd.testing.assert_frame_equal( df1.reset_index(drop=True), df2.reset_index(drop=True), check_like=True, ) def test_aggregation_missing_file_returns_empty(migrate_module, tmp_path): """文件不存在 -> 空 DataFrame(不崩)""" hist = tmp_path / "hist" hist.mkdir() df = migrate_module._read_announce_union_aggregated(str(hist), ["000852"]) assert df.empty assert list(df.columns) == [ "index_code", "code", "code_name", "in_current", "was_removed", "source"] def test_aggregation_prefers_snapshot_name(migrate_module, tmp_path): """code_name 优先 snapshot 当前名(announce 的历史名可能过时)""" ann = _make_announce_df([ ("000001", "旧名", "add"), ("000001", "旧名", "current"), ]) snap = pd.DataFrame({ "index_code": ["000852"], "code": ["000001"], "code_name": ["新名"], }) hist = tmp_path / "hist" hist.mkdir() ann.to_parquet(hist / "000852_announce_union.parquet", index=False) snap.to_parquet(hist / "000852_snapshot.parquet", index=False) df = migrate_module._read_announce_union_aggregated(str(hist), ["000852"]) assert df.iloc[0]["code_name"] == "新名" # ============================================================ # Integration: 真实 announce_union.parquet (Mac parse 产出) # ============================================================ MAC_HIST = ROOT / "data" / "index_const_hist" @pytest.mark.integration def test_real_000852_distinct_greater_1000(migrate_module): """000852 治偏差证据: distinct > 1000(plan 验证标准)""" if not (MAC_HIST / "000852_announce_union.parquet").exists(): pytest.skip("000852_announce_union.parquet 未生成, 先跑 parse_csindex_announce.py") df = migrate_module._read_announce_union_aggregated(str(MAC_HIST), ["000852"]) n_distinct = df["code"].nunique() n_removed = (df["was_removed"] == 1).sum() assert n_distinct > 1000, f"000852 distinct={n_distinct}, 期望 >1000(治偏差)" assert n_removed > 0, f"000852 was_removed={n_removed}, 期望 >0" print(f"\n000852: distinct={n_distinct}, in_current={int((df['in_current']==1).sum())}, " f"was_removed={n_removed}") @pytest.mark.integration def test_real_932000_distinct_approx_2000(migrate_module): """932000 launch xlsx bug 修复: distinct ≈ 2000+(launch ∪ current)""" if not (MAC_HIST / "932000_announce_union.parquet").exists(): pytest.skip("932000_announce_union.parquet 未生成, 先跑 parse_csindex_announce.py") df = migrate_module._read_announce_union_aggregated(str(MAC_HIST), ["932000"]) n_distinct = df["code"].nunique() # launch 修复后 distinct 不再=1, 应在 2000+(launch ∪ current 重叠后 2684 实证) assert n_distinct >= 2000, f"932000 distinct={n_distinct}, 期望 ≥2000(launch 修复)" print(f"\n932000: distinct={n_distinct}, in_current={int((df['in_current']==1).sum())}, " f"was_removed={int((df['was_removed']==1).sum())}")