live-hub-py/tests/test_login_attempt_migration.py

"""迁移 20261001_0035 的存量数据归一验证（在真实 MySQL 临时库上执行）。

线上升级时要保证：
- 持有 CK 的记录 status 归一到 success，历史批次状态搬进 attempt_*；
- 没有 CK 且卡在 pending/running 的记录收敛为 failed；
- 任何情况下都不丢 CK。
"""

import secrets
from pathlib import Path
from urllib.parse import quote_plus

from alembic import command
from alembic.config import Config
from sqlalchemy import create_engine, text

import web.backend.database as database_module
from tests.support.db import get_mysql_credentials, get_test_engine

PROJECT_ROOT = Path(__file__).resolve().parents[1]


def _migration_config(url: str) -> Config:
    config = Config(str(PROJECT_ROOT / "alembic.ini"))
    config.set_main_option(
        "script_location", str(PROJECT_ROOT / "web" / "backend" / "migrations")
    )
    config.set_main_option("sqlalchemy.url", url)
    return config


def _insert_task(conn, *, status: str, cookie: str, message: str = "") -> None:
    conn.execute(
        text(
            """
            INSERT INTO login_tasks
                (batch_id, account_id, status, cookie, message, created_by,
                 created_at, finished_at)
            VALUES
                ('batch-1', 1, :status, :cookie, :message, 1,
                 '2026-09-30 00:00:00', '2026-09-30 01:00:00')
            """
        ),
        {"status": status, "cookie": cookie, "message": message},
    )


class TestLoginAttemptMigration:
    def setup_method(self):
        creds = get_mysql_credentials()
        self.mig_db = f"douyu_login_mig_{secrets.token_hex(4)}"
        self.url = (
            f"mysql+pymysql://{quote_plus(creds['user'])}:{quote_plus(creds['password'])}"
            f"@{creds['host']}:{creds['port']}/{self.mig_db}?charset=utf8mb4"
        )
        self.admin_engine = get_test_engine()
        with self.admin_engine.begin() as conn:
            conn.execute(
                text(
                    f"CREATE DATABASE `{self.mig_db}` CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci;"
                )
            )

        self.database_module = database_module
        self.original_url = database_module.DATABASE_URL
        # env.py 用模块级 DATABASE_URL 覆盖配置里的连接串。
        database_module.DATABASE_URL = self.url
        self.config = _migration_config(self.url)
        command.upgrade(self.config, "20260901_0034")

        self.engine = create_engine(self.url)
        with self.engine.begin() as conn:
            conn.execute(
                text(
                    """
                    INSERT INTO accounts
                        (id, username, password, email, email_password,
                         email_imap_server, email_imap_port, email_imap_ssl,
                         tag, remark, uid, nickname, created_at, updated_at)
                    VALUES
                        (1, '远行者103', 'p', 'e@example.com', 'm',
                         '', 993, 1, '', '', '', '', '2026-09-30 00:00:00',
                         '2026-09-30 00:00:00')
                    """
                )
            )
            conn.execute(
                text(
                    """
                    INSERT INTO users (id, username, password_hash, role, is_active)
                    VALUES (1, 'admin', 'hash', 'super_admin', 1)
                    """
                )
            )
            _insert_task(conn, status="success", cookie="ck-success")
            _insert_task(
                conn,
                status="relogin_failed",
                cookie="ck-relogin",
                message="重新登录失败: 密码错误（旧 Cookie 已保留）",
            )
            _insert_task(conn, status="pending", cookie="ck-pending")
            _insert_task(conn, status="pending", cookie="")
            _insert_task(conn, status="running", cookie="")
            _insert_task(conn, status="failed", cookie="ck-failed", message="登录失败")

    def teardown_method(self):
        self.engine.dispose()
        self.database_module.DATABASE_URL = self.original_url
        with self.admin_engine.begin() as conn:
            conn.execute(text(f"DROP DATABASE IF EXISTS `{self.mig_db}`;"))

    def _rows(self) -> list[dict]:
        with self.engine.connect() as conn:
            rows = conn.execute(
                text(
                    """
                    SELECT status, cookie, message, attempt_status,
                           attempt_message, attempt_batch_id, attempt_at
                      FROM login_tasks
                     ORDER BY id
                    """
                )
            ).mappings()
            return [dict(row) for row in rows]

    def test_upgrade_normalizes_existing_rows(self):
        command.upgrade(self.config, "head")

        rows = self._rows()
        (
            success,
            relogin_failed,
            pending_with_ck,
            pending_without_ck,
            running_without_ck,
            failed_with_ck,
        ) = rows

        # 已经成功的记录：status 不变，也没有历史尝试需要搬迁
        assert success["status"] == "success"
        assert success["cookie"] == "ck-success"
        assert success["attempt_status"] in (None, "")

        # 重登失败：CK 保留，status 归一为 success，失败信息进 attempt_*
        assert relogin_failed["cookie"] == "ck-relogin"
        assert relogin_failed["status"] == "success"
        assert relogin_failed["attempt_status"] == "failed"
        assert "旧 Cookie 已保留" in relogin_failed["attempt_message"]
        assert relogin_failed["attempt_batch_id"] == "batch-1"

        # 卡在 pending 但持有 CK：同样归一为 success + 中断说明
        assert pending_with_ck["cookie"] == "ck-pending"
        assert pending_with_ck["status"] == "success"
        assert pending_with_ck["attempt_status"] == "failed"
        assert "旧 Cookie 已保留" in pending_with_ck["attempt_message"]

        # 没有 CK 的在途记录：收敛为失败，不再永久排队
        for row in (pending_without_ck, running_without_ck):
            assert row["status"] == "failed"
            assert row["attempt_status"] == "failed"
            assert "服务重启" in row["message"]

        # 复用记录上的旧 CK 不能被清掉
        assert failed_with_ck["status"] == "success"
        assert failed_with_ck["cookie"] == "ck-failed"
        assert failed_with_ck["attempt_status"] == "failed"