"""成份股 merge / migrate pipeline 可重跑(idempotent)回归测试。 背景: 方案A 落地后 VPS 上的状态是 - ``constituent_unified`` 已存在(正式表) - ``bs_index_constituent_old`` 已存在(baostock 300/500/50 全部历史时点, 权威历史源) - ``bs_index_constituent`` 不存在(被 rename 走了) - ``constituent_unified_staging`` 仅在 migrate 跑完后存在 旧 ``merge_constituent.py`` 用 RENAME + 旧 ``migrate_constituent.py`` 直接 ``FROM bs_index_constituent`` → 月度 schtask 再跑会崩。这组测试断言: 1. ``merge`` 跑两次都不崩, 行数稳定(DROP+CREATE from staging 幂等, 不用 RENAME), 且不碰 ``bs_index_constituent_old``。 2. ``migrate`` 跑两次都不崩, 从 ``bs_index_constituent_old`` 读, staging 行数稳定。 全部在 Mac 本地用 tmp sqlite + 空 HIST 目录跑, 不依赖 VPS / baostock / akshare。 """ from __future__ import annotations import os import sqlite3 import sys import pytest # 让测试能 import scripts/data_platform/ 下的模块 _HERE = os.path.dirname(os.path.abspath(__file__)) _SCRIPT_DIR = os.path.abspath(os.path.join(_HERE, "..", "..", "scripts", "data_platform")) if _SCRIPT_DIR not in sys.path: sys.path.insert(0, _SCRIPT_DIR) import merge_constituent as mc # noqa: E402 import migrate_constituent as mig # noqa: E402 # ======================== helpers ======================== def _create_staging(db_path: str, rows: list[tuple[str, str, str, str, int, int]]) -> None: """在 tmp db 里造一张 constituent_unified_staging。 rows: (index_code, code, code_name, source, in_current, was_removed) """ with sqlite3.connect(db_path) as c: c.execute("PRAGMA busy_timeout = 60000") c.execute( "CREATE TABLE constituent_unified_staging (" "index_code TEXT, code TEXT, code_name TEXT, source TEXT, " "in_current INTEGER, was_removed INTEGER)" ) c.executemany( "INSERT INTO constituent_unified_staging VALUES (?,?,?,?,?,?)", rows ) c.commit() def _create_bs_old(db_path: str, rows: list[tuple[str, str, str, str]]) -> None: """在 tmp db 里造一张 bs_index_constituent_old(updateDate, index_code, code, code_name)。 默认 index_code 用 baostock 原始命名(hs300/zz500/sz50), migrate 内有 BS_MAP 映射。 """ with sqlite3.connect(db_path) as c: c.execute("PRAGMA busy_timeout = 60000") c.execute( "CREATE TABLE bs_index_constituent_old (" "updateDate TEXT, index_code TEXT, code TEXT, code_name TEXT)" ) c.executemany( "INSERT INTO bs_index_constituent_old VALUES (?,?,?,?)", rows ) c.commit() def _table_exists(db_path: str, table: str) -> bool: with sqlite3.connect(db_path) as c: row = c.execute( "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?", (table,) ).fetchone() return row is not None def _count_rows(db_path: str, table: str) -> int: with sqlite3.connect(db_path) as c: return c.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0] # ======================== test 1: merge 幂等 ======================== def test_merge_idempotent_rerun(tmp_path): """merge 跑两次不崩, constituent_unified 行数 == staging 行数; _old 不被碰。 断言的是 DROP+CREATE from staging 语义(幂等), 而不是 RENAME(只能跑一次)。 """ db = tmp_path / "test.db" db_path = str(db) # staging 5 行 (300 + 500 各几只) staging_rows = [ ("000300", "600519", "贵州茅台", "baostock", 1, 0), ("000300", "601318", "中国平安", "baostock", 1, 0), ("000300", "000001", "平安银行", "baostock", 0, 1), ("000905", "002415", "海康威视", "baostock", 1, 0), ("000905", "300750", "宁德时代", "baostock", 1, 0), ] _create_staging(db_path, staging_rows) # 模拟方案A后状态: bs_index_constituent_old 已存在(不该被 merge 碰) _create_bs_old(db_path, [ ("2024-06-17", "hs300", "sh.600519", "贵州茅台"), ("2024-06-17", "hs300", "sh.601318", "中国平安"), ]) bs_old_rows_before = _count_rows(db_path, "bs_index_constituent_old") # Act 1: 第一次 merge mc.merge(db_path) # Assert 1: constituent_unified 存在, 行数 == staging 行数 assert _table_exists(db_path, "constituent_unified") is True assert _count_rows(db_path, "constituent_unified") == len(staging_rows) # staging 仍可读(没被 rename 走) assert _count_rows(db_path, "constituent_unified_staging") == len(staging_rows) # bs_old 没被碰 assert _count_rows(db_path, "bs_index_constituent_old") == bs_old_rows_before # Act 2: 第二次 merge —— 关键回归(RENAME 方案这里会崩) mc.merge(db_path) # Assert 2: 行数稳定, 没异常 assert _count_rows(db_path, "constituent_unified") == len(staging_rows) assert _count_rows(db_path, "constituent_unified_staging") == len(staging_rows) assert _count_rows(db_path, "bs_index_constituent_old") == bs_old_rows_before # ======================== test 2: migrate 幂等 ======================== def test_migrate_idempotent_rerun_from_old(tmp_path, monkeypatch): """migrate 跑两次不崩, 从 bs_index_constituent_old 读, staging 行数稳定。 构造 baostock 原始 index_code (hs300/zz500/sz50), migrate 内 BS_MAP 映射到 000300/000905/000016。HIST 目录用空目录(深证/中证 parquet 缺失, migrate 应优雅产出 0 行不崩)。 """ db = tmp_path / "test.db" db_path = str(db) # baostock 历史时点(2 个日期, 模拟一只股被踢出) bs_rows = [ # 2024-06-17 时点: hs300 含 600519 / 601318 ("2024-06-17", "hs300", "sh.600519", "贵州茅台"), ("2024-06-17", "hs300", "sh.601318", "中国平安"), # 2024-12-16 时点: hs300 600519 留下, 601318 被踢, 新进 000001 ("2024-12-16", "hs300", "sh.600519", "贵州茅台"), ("2024-12-16", "hs300", "sz.000001", "平安银行"), # zz500 一只 ("2024-06-17", "zz500", "sh.600036", "招商银行"), # sz50 一只 ("2024-06-17", "sz50", "sh.600000", "浦发银行"), ] _create_bs_old(db_path, bs_rows) # HIST 指向空目录(深证 union / 中证 snapshot 都缺失, migrate 应优雅跳过) empty_hist = tmp_path / "index_const_hist" empty_hist.mkdir() monkeypatch.setattr(mig, "HIST", str(empty_hist)) # Act 1 mig.migrate(db_path) # Assert 1: staging 建好, 行数 == 去重后的 _old 行数 assert _table_exists(db_path, "constituent_unified_staging") is True # 去重后 6 个唯一 (index_code, code) 对: # hs300: 600519 / 601318 / 000001 (3) # zz500: 600036 (1) # sz50: 600000 (1) # 注意 600519 在两个时点出现, groupby first() 去重为 1 行 → 共 5 行 expected_rows = 5 assert _count_rows(db_path, "constituent_unified_staging") == expected_rows # 内容抽查: hs300 映射对了, was_removed 标记对了 with sqlite3.connect(db_path) as c: # hs300 应有 3 行 n_hs300 = c.execute( "SELECT COUNT(*) FROM constituent_unified_staging WHERE index_code='000300'" ).fetchone()[0] assert n_hs300 == 3 # 601318 在最后时点已不在 → was_removed=1 row_601318 = c.execute( "SELECT in_current, was_removed FROM constituent_unified_staging " "WHERE index_code='000300' AND code='601318'" ).fetchone() assert row_601318 == (0, 1) # 600519 仍在 → in_current=1 row_600519 = c.execute( "SELECT in_current, was_removed FROM constituent_unified_staging " "WHERE index_code='000300' AND code='600519'" ).fetchone() assert row_600519 == (1, 0) # Act 2: 第二次 migrate —— 关键回归 mig.migrate(db_path) # Assert 2: staging 行数稳定(DROP+rebuild 幂等) assert _count_rows(db_path, "constituent_unified_staging") == expected_rows # _old 没被碰 assert _count_rows(db_path, "bs_index_constituent_old") == len(bs_rows) def test_migrate_reads_bs_old_not_live_when_both_exist(tmp_path, monkeypatch): """migrate 兼容性: _old 和 live 表都在时, 两边 UNION 不丢数据(robustness)。 方案A 后正常只有 _old; 但若将来有人重建了 live bs_index_constituent, migrate 应把两者都读(UNION ALL 去重), 兼容两种状态。 """ db = tmp_path / "test.db" db_path = str(db) # _old 里有 hs300 600519 _create_bs_old(db_path, [ ("2024-06-17", "hs300", "sh.600519", "贵州茅台"), ]) # live 表里补一只新的(模拟未来重建 live 后采到新时点) with sqlite3.connect(db_path) as c: c.execute( "CREATE TABLE bs_index_constituent (" "updateDate TEXT, index_code TEXT, code TEXT, code_name TEXT)" ) c.execute( "INSERT INTO bs_index_constituent VALUES (?,?,?,?)", ("2025-06-16", "hs300", "sh.688981", "中芯国际"), ) c.commit() empty_hist = tmp_path / "index_const_hist" empty_hist.mkdir() monkeypatch.setattr(mig, "HIST", str(empty_hist)) # Act mig.migrate(db_path) # Assert: staging 应同时含 _old 的 600519 + live 的 688981 (UNION 后去重 2 行) with sqlite3.connect(db_path) as c: codes = sorted( r[0] for r in c.execute( "SELECT code FROM constituent_unified_staging " "WHERE index_code='000300'" ).fetchall() ) assert codes == ["600519", "688981"]