PayrollManagementSystem/scripts/run_isolated_tests.py

"""使用唯一的本地测试库运行迁移和集成测试，绝不清理原业务库。"""

import argparse
import os
import re
import secrets
import subprocess
import sys
from pathlib import Path


def run_mysql(sql: str) -> subprocess.CompletedProcess:
    """密码仅在容器环境和标准输入中使用，不回显 SQL 错误内容。"""
    return subprocess.run(
        [
            "docker",
            "exec",
            "-i",
            "ys-mysql-dev",
            "sh",
            "-c",
            'export MYSQL_PWD="$MYSQL_ROOT_PASSWORD"; exec mysql -uroot',
        ],
        input=sql.encode(),
        capture_output=True,
    )


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--reuse-database")
    parser.add_argument("--migration-roundtrip", action="store_true")
    parser.add_argument("--keep", action="store_true", help="保留本次隔离库和账号供复查")
    parsed, test_args = parser.parse_known_args()
    if test_args[:1] == ["--"]:
        test_args = test_args[1:]
    if parsed.reuse_database and not re.fullmatch(
        r"payroll_test_[0-9a-f]{12}", parsed.reuse_database
    ):
        raise SystemExit("只允许复用本运行器创建的隔离测试库")
    root = Path(__file__).resolve().parents[1]
    python = Path(sys.executable)
    if sys.prefix == sys.base_prefix:
        for candidate in (".venv", ".venv-workbench"):
            local_python = root / "backend" / candidate / "Scripts/python.exe"
            if local_python.exists():
                python = local_python
                break
    suffix = secrets.token_hex(6)
    database = parsed.reuse_database or "payroll_test_" + suffix
    username = "pytest_" + suffix
    password = secrets.token_urlsafe(30)
    created_database = created_user = False
    result_code = 1
    try:
        if not parsed.reuse_database:
            if run_mysql(
                f"CREATE DATABASE `{database}` CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci;"
            ).returncode:
                raise RuntimeError("隔离库创建失败，未执行迁移或测试")
            created_database = True
        if run_mysql(f"CREATE USER '{username}'@'%' IDENTIFIED BY '{password}';").returncode:
            raise RuntimeError("隔离账号创建失败，未执行迁移或测试")
        created_user = True
        if run_mysql(f"GRANT ALL PRIVILEGES ON `{database}`.* TO '{username}'@'%';").returncode:
            raise RuntimeError("隔离账号授权失败，未执行迁移或测试")
        env = os.environ | {
            "MYSQL_DATABASE": database,
            "MYSQL_USER": username,
            "MYSQL_PASSWORD": password,
            "PAYROLL_TEST_DATABASE": database,
            "REDIS_DB": "15",
            "PYTHONIOENCODING": "utf-8",
            "UPLOAD_DIR": str(root / ".ai" / "isolated-uploads" / database),
        }
        print("独立测试库：" + database, flush=True)
        commands = [["-m", "alembic", "upgrade", "head"]]
        if parsed.migration_roundtrip:
            commands += [
                ["-m", "alembic", "downgrade", "07b8c9d0e1f2"],
                ["-m", "alembic", "upgrade", "head"],
            ]
        commands.append(["-m", "pytest", *(test_args or ["tests/integration"]), "-q", "--tb=short"])
        for args in commands:
            result_code = subprocess.run(
                [str(python), *args], cwd=root / "backend", env=env
            ).returncode
            if result_code:
                break
    except RuntimeError as exc:
        print(str(exc), file=sys.stderr)
    finally:
        # 只清理本次 CREATE 已成功的对象；重名创建失败及复用库绝不能被删除。
        if not parsed.keep:
            cleanup = []
            if created_database:
                cleanup.append(f"DROP DATABASE `{database}`;")
            if created_user:
                cleanup.append(f"DROP USER '{username}'@'%';")
            for statement in cleanup:
                if run_mysql(statement).returncode:
                    print(f"隔离测试资源清理失败，请核对 {database} / {username}", file=sys.stderr)
                    result_code = result_code or 1
    return result_code


if __name__ == "__main__":
    raise SystemExit(main())