# 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)