re-xxm/tests/test_migrations.py

# -*- coding: utf-8 -*-
"""SQLite 迁移测试（不联网）: uv run python tests/test_migrations.py"""

import os
import sqlite3
import sys
import tempfile
from pathlib import Path

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.stdout.reconfigure(encoding="utf-8")

from reaction_board.migrations import (  # noqa: E402
    MIGRATIONS,
    Migration,
    MigrationError,
    current_version,
    prune_backups,
    run_migrations,
)

EXPECTED_TABLES = {
    "video", "snapshot", "folder", "folder_resource", "up_blacklist",
    "up_stat", "up_profile", "candidate", "seen_aid", "notify", "fav_action", "kv",
}


def table_names(conn):
    rows = conn.execute(
        "SELECT name FROM sqlite_master WHERE type = 'table' AND name NOT LIKE 'sqlite_%'"
    ).fetchall()
    return {row[0] for row in rows}


def test_fresh_and_idempotent():
    with tempfile.TemporaryDirectory() as tmp:
        conn = sqlite3.connect(Path(tmp) / "app.db")
        try:
            applied = run_migrations(conn)
            assert applied == [1, 2, 3, 4]
            assert current_version(conn) == 4
            assert EXPECTED_TABLES <= table_names(conn)
            assert "schema_version" in table_names(conn)
            columns = {
                row[1] for row in conn.execute("PRAGMA table_info(folder_resource)")
            }
            assert "fav_time" in columns
            candidate_columns = {
                row[1] for row in conn.execute("PRAGMA table_info(candidate)")
            }
            assert {"duration", "play", "like", "favorite", "danmaku", "review"} <= candidate_columns
            video_columns = {
                row[1] for row in conn.execute("PRAGMA table_info(video)")
            }
            assert {
                "pages", "duration", "up_follower", "up_follower_at",
                "up_follower_backfilled",
            } <= video_columns
            profile_columns = {
                row[1] for row in conn.execute("PRAGMA table_info(up_profile)")
            }
            assert {
                "mid", "follower_first", "following_first", "first_captured_at",
                "follower_latest", "last_refreshed",
            } <= profile_columns
            # 再跑一次不应重复应用
            assert run_migrations(conn) == []
            row = conn.execute("SELECT COUNT(*) FROM schema_version").fetchone()
            assert row[0] == 4
        finally:
            conn.close()


def test_backup_before_upgrade():
    with tempfile.TemporaryDirectory() as tmp:
        db_path = Path(tmp) / "app.db"
        backup_dir = Path(tmp) / "backup"
        conn = sqlite3.connect(db_path)
        try:
            run_migrations(conn, backup_dir=backup_dir)
            conn.execute("INSERT INTO kv(k, v) VALUES ('hello', 'world')")
            conn.commit()

            v5 = Migration(5, "add probe", ("CREATE TABLE probe (id INTEGER PRIMARY KEY)",))
            applied = run_migrations(
                conn,
                backup_dir=backup_dir,
                now=1_800_000_000,
                migrations=MIGRATIONS + (v5,),
            )
            assert applied == [5]
            backups = list(backup_dir.glob("app-pre-v4-*.db"))
            assert len(backups) == 1
            check = sqlite3.connect(backups[0])
            try:
                assert check.execute("SELECT v FROM kv WHERE k='hello'").fetchone()[0] == "world"
                assert "probe" not in table_names(check)
            finally:
                check.close()
        finally:
            conn.close()


def test_failed_migration_rolls_back():
    with tempfile.TemporaryDirectory() as tmp:
        conn = sqlite3.connect(Path(tmp) / "app.db")
        bad = Migration(
            1,
            "broken",
            (
                "CREATE TABLE should_not_exist (id INTEGER PRIMARY KEY)",
                "CREATE TABLE broken (",
            ),
        )
        try:
            try:
                run_migrations(conn, migrations=(bad,))
            except MigrationError as exc:
                assert "已回滚" in str(exc)
            else:
                raise AssertionError("坏迁移应抛错")
            assert "should_not_exist" not in table_names(conn)
            assert current_version(conn) == 0
        finally:
            conn.close()


def test_prune_backups():
    with tempfile.TemporaryDirectory() as tmp:
        backup_dir = Path(tmp)
        old = backup_dir / "app-pre-v1-old.db"
        fresh = backup_dir / "app-pre-v1-fresh.db"
        other = backup_dir / "keep.txt"
        for path in (old, fresh, other):
            path.write_bytes(b"x")
        now = 1_800_000_000
        old_time = now - 20 * 86400
        os.utime(old, (old_time, old_time))
        os.utime(fresh, (now, now))
        removed = prune_backups(backup_dir, keep_days=14, now=now)
        assert removed == [old]
        assert fresh.exists() and other.exists()


def test_validate_migration_order():
    with tempfile.TemporaryDirectory() as tmp:
        conn = sqlite3.connect(Path(tmp) / "app.db")
        try:
            try:
                run_migrations(conn, migrations=(
                    Migration(2, "x", ()),
                    Migration(1, "y", ()),
                ))
            except ValueError as exc:
                assert "升序" in str(exc)
            else:
                raise AssertionError("乱序迁移应失败")
        finally:
            conn.close()


def main():
    test_fresh_and_idempotent()
    test_backup_before_upgrade()
    test_failed_migration_rolls_back()
    test_prune_backups()
    test_validate_migration_order()
    print("OK: SQLite 迁移、备份与回滚测试全部通过")


if __name__ == "__main__":
    main()