FinDataHub/Scripts/bin/tests/test_archive_data.py

# coding=utf-8
"""archive_data.py 数据归档脚本单元测试.

重点覆盖：
- ensure_archive_table：归档表存在判断与 CREATE TABLE LIKE
- archive_table：count=0 短路、dry 只统计、正常迁移+删除
- main：表存在过滤、--table 指定、commit/OPTIMIZE 时机
"""

import json
import os
import sys
from collections import deque
from types import SimpleNamespace
from unittest.mock import MagicMock

import pytest

# 将 Scripts/bin 加入 sys.path
_bin_dir = os.path.join(os.path.dirname(__file__), "..")
if _bin_dir not in sys.path:
    sys.path.insert(0, _bin_dir)

import archive_data

# conftest 注入的 pymysql mock 缺少 cursors，补齐（db_pool 导入可能需要）
_pymysql_mock = sys.modules.get("pymysql")
if _pymysql_mock is not None and not hasattr(_pymysql_mock, "cursors"):
    _pymysql_mock.cursors = SimpleNamespace(DictCursor=dict)


class _FakeCursor:
    """记录 execute，按队列返回 fetchone 结果."""

    def __init__(self, fetchone_results=None):
        self.executed = []
        self.rowcount = 0
        self._results = deque(fetchone_results or [])
        self.closed = False

    def execute(self, sql, params=None):
        self.executed.append((sql, params or ()))

    def fetchone(self):
        return self._results.popleft() if self._results else None

    def close(self):
        self.closed = True


def _executed_sql(cursor):
    return "\n".join(s for s, _ in cursor.executed)


# ── ensure_archive_table ────────────────────────────────────


class TestEnsureArchiveTable:
    def test_returns_existing_archive(self):
        cur = _FakeCursor([(1,)])
        assert archive_data.ensure_archive_table(cur, "news_flash") == "news_flash_archive"
        # 未执行 CREATE
        assert "CREATE TABLE" not in _executed_sql(cur)
        assert "information_schema" in _executed_sql(cur)

    def test_creates_when_missing(self, capsys):
        cur = _FakeCursor([None])
        assert archive_data.ensure_archive_table(cur, "news_cls") == "news_cls_archive"
        sql = _executed_sql(cur)
        assert "CREATE TABLE `news_cls_archive` LIKE `news_cls`" in sql
        assert "创建归档表" in capsys.readouterr().out


# ── archive_table ───────────────────────────────────────────


class TestArchiveTable:
    def test_zero_count_returns_early(self):
        # archive 表已存在 (1,)，count=0 (0,)
        cur = _FakeCursor([(1,), (0,)])
        assert archive_data.archive_table(cur, "news_flash", 1000) == (0, 0)
        sql = _executed_sql(cur)
        assert "INSERT INTO" not in sql
        assert "DELETE FROM" not in sql

    def test_dry_run_counts_only(self, capsys):
        # archive 已存在 (1,)，count=5
        cur = _FakeCursor([(1,), (5,)])
        assert archive_data.archive_table(cur, "news_flash", 1000, dry=True) == (5, 0)
        sql = _executed_sql(cur)
        assert "INSERT INTO" not in sql
        assert "DELETE FROM" not in sql
        assert "待归档 5 条" in capsys.readouterr().out

    def test_normal_migrate_and_delete(self):
        # archive 已存在 (1,)，count=10；INSERT rowcount=10，DELETE rowcount=10
        cur = _FakeCursor([(1,), (10,)])
        cur.rowcount = 10
        assert archive_data.archive_table(cur, "news_flash", 1000) == (10, 0)
        sql = _executed_sql(cur)
        assert "INSERT INTO `news_flash_archive` SELECT * FROM `news_flash`" in sql
        assert "DELETE FROM `news_flash` WHERE published_at < %s" in sql

    def test_creates_archive_table_then_migrates(self):
        # archive 不存在 (None)，count=5
        cur = _FakeCursor([None, (5,)])
        cur.rowcount = 5
        assert archive_data.archive_table(cur, "news_hotlist", 1000) == (5, 0)
        sql = _executed_sql(cur)
        assert "CREATE TABLE `news_hotlist_archive`" in sql
        assert "INSERT INTO" in sql


# ── main ────────────────────────────────────────────────────


class _FakeConn:
    def __init__(self, cursor, ctx):
        self._cursor = cursor
        self._ctx = ctx

    def cursor(self):
        return self._cursor

    def commit(self):
        self._ctx.committed = True


class _FakeCtx:
    """模拟 db_pool.DBContext 上下文管理器."""

    def __init__(self, cursor):
        self._cursor = cursor
        self.committed = False

    def __enter__(self):
        self.conn = _FakeConn(self._cursor, self)
        return self.conn

    def __exit__(self, *a):
        return False


def _run_main(monkeypatch, argv, cursor):
    ctx = _FakeCtx(cursor)
    monkeypatch.setattr(archive_data, "DBContext", lambda: ctx)
    monkeypatch.setattr(sys, "argv", ["archive_data.py"] + argv)
    archive_data.main()
    return ctx


class TestMain:
    def test_dry_run_single_table(self, monkeypatch, capsys):
        # 表存在 (1,)，archive 存在 (1,)，count=5 → dry 不迁移不 commit
        cur = _FakeCursor([(1,), (1,), (5,)])
        ctx = _run_main(monkeypatch, ["--days", "30", "--dry", "--table", "news_flash"], cur)
        assert ctx.committed is False
        sql = _executed_sql(cur)
        assert "INSERT INTO" not in sql
        assert "OPTIMIZE" not in sql
        assert "共 5 条" in capsys.readouterr().out

    def test_missing_table_skipped(self, monkeypatch, capsys):
        # 表不存在 → continue，total=0
        cur = _FakeCursor([None])
        _run_main(monkeypatch, ["--table", "news_missing"], cur)
        assert "共 0 条" in capsys.readouterr().out

    def test_normal_commits_and_optimizes(self, monkeypatch):
        # 表存在 (1,)，archive 存在 (1,)，count=5，OPTIMIZE 前再查存在 (1,)
        cur = _FakeCursor([(1,), (1,), (5,), (1,)])
        cur.rowcount = 5
        ctx = _run_main(monkeypatch, ["--table", "news_flash"], cur)
        assert ctx.committed is True
        sql = _executed_sql(cur)
        assert "INSERT INTO `news_flash_archive`" in sql
        assert "DELETE FROM `news_flash`" in sql
        assert "OPTIMIZE TABLE `news_flash`" in sql

    def test_all_tables_when_no_filter(self, monkeypatch, capsys):
        # 遍历 ARCHIVE_TABLES，全部存在且 count=0
        cur = _FakeCursor([])
        _run_main(monkeypatch, ["--dry"], cur)
        # 表名通过参数传递（参数化查询），从 params 收集
        checked = set(p[0] for _, p in cur.executed if p and p[0] in archive_data.ARCHIVE_TABLES)
        assert checked == set(archive_data.ARCHIVE_TABLES)