sanguo_vnpy_v2/tests/portfolio/test_merge_constituent_rerun.py

"""成份股 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"]