group_fqcd_jr/tools/run_tests_on_test_db.py

"""在独立测试库上跑测试：与开发用的 worker 进程物理隔离。

为什么需要：`python -m app.worker` 会认领任意 `queued/running/cancel_requested` 运行，
并消费 Outbox 事件。只要它与测试共用同一个库，测试造的数据就会被并发消费，表现为
"同一个用例全量跑失败、单跑通过"的偶发失败——那不是代码缺陷，是环境干扰。

流程（每一步都只作用于测试库）：

1. 由 `.env` 的 `MYSQL_DSN` 推导测试库 DSN（只换库名，其余原样）；
2. 对测试库执行 `alembic upgrade head`（合并两条分支后，迁移链已能从空库建出全部表，
   不再需要手工导入 `alembic/baseline_generated.sql`）；`--rebuild` 会先清空测试库再重建；
3. 灌入 RBAC 测试账号（`tools/seed_test_rbac.py`，9001/9002/9003 号段）；
4. 仅在本进程内用 `MYSQL_DSN` 覆盖，跑 pytest。

前置（只需一次，用有建库权限的账号执行）：

    CREATE DATABASE IF NOT EXISTS jr_agent_test
      CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci;
    GRANT ALL PRIVILEGES ON jr_agent_test.* TO 'jr_app'@'localhost';

用法：

    python -m tools.run_tests_on_test_db                 # 默认跑 tests/integration
    python -m tools.run_tests_on_test_db tests           # 跑全部测试目录
    python -m tools.run_tests_on_test_db --rebuild tests # 先清空测试库再由迁移重建
"""

from __future__ import annotations

import os
import subprocess
import sys
from pathlib import Path

import pymysql
from sqlalchemy.engine import make_url

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from app.core.config import get_settings  # noqa: E402

DEFAULT_DATABASE = "jr_agent_test"


def _connect(**overrides: object) -> pymysql.connections.Connection:
    settings = get_settings()
    url = make_url(settings.mysql_dsn)
    return pymysql.connect(
        host=url.host or "127.0.0.1",
        port=url.port or 3306,
        user=url.username,
        password=url.password or "",
        charset="utf8mb4",
        **overrides,  # type: ignore[arg-type]
    )


def _table_count(connection: pymysql.connections.Connection, database: str) -> int:
    with connection.cursor() as cursor:
        cursor.execute(
            "SELECT COUNT(*) FROM information_schema.TABLES WHERE TABLE_SCHEMA = %s",
            (database,),
        )
        return int(cursor.fetchone()[0])


def _drop_all_tables(connection: pymysql.connections.Connection, database: str) -> int:
    """清空测试库（仅测试库），返回删除的表数。

    先关外键检查再逐张 DROP：表之间存在外键，顺序删除会互相阻塞。
    """
    with connection.cursor() as cursor:
        cursor.execute(f"USE `{database}`")
        cursor.execute(
            "SELECT table_name FROM information_schema.tables WHERE table_schema=%s",
            (database,),
        )
        tables = [row[0] for row in cursor.fetchall()]
        cursor.execute("SET FOREIGN_KEY_CHECKS=0")
        for table in tables:
            cursor.execute(f"DROP TABLE IF EXISTS `{table}`")
        cursor.execute("SET FOREIGN_KEY_CHECKS=1")
    connection.commit()
    return len(tables)


def _run(step: str, args: list[str], env: dict[str, str]) -> None:
    print(f"\n>>> {step}: {' '.join(args)}", flush=True)
    completed = subprocess.run(args, cwd=ROOT, env=env, check=False)
    if completed.returncode != 0:
        raise SystemExit(f"{step} 失败（退出码 {completed.returncode}）")


def main() -> None:
    database = os.environ.get("TEST_DB_NAME", DEFAULT_DATABASE)
    arguments = sys.argv[1:]
    rebuild = "--rebuild" in arguments
    pytest_args = [item for item in arguments if item != "--rebuild"] or ["tests/integration"]

    base_url = make_url(get_settings().mysql_dsn)
    test_url = base_url.set(database=database)
    if base_url.database == database:
        raise SystemExit(f"测试库不能与开发库同名（当前：{database}）")

    try:
        connection = _connect(database=database)
    except pymysql.err.OperationalError as exc:
        raise SystemExit(
            f"无法连接测试库 {database}：{exc}\n"
            "请先用有建库权限的账号执行：\n"
            f"  CREATE DATABASE IF NOT EXISTS {database} "
            "CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci;\n"
            f"  GRANT ALL PRIVILEGES ON {database}.* TO 'jr_app'@'localhost';"
        ) from exc

    with connection:
        tables = _table_count(connection, database)
        print(f"测试库 {database} 现有表：{tables}")
        if rebuild and tables:
            dropped = _drop_all_tables(connection, database)
            print(f"--rebuild：已清空测试库（删除 {dropped} 张表）")

    env = {
        **os.environ,
        "MYSQL_DSN": test_url.render_as_string(hide_password=False),
    }
    python = sys.executable
    _run("结构迁移", [python, "-m", "alembic", "upgrade", "head"], env)
    _run("灌入测试账号", [python, "-m", "tools.seed_test_rbac"], env)
    _run("执行测试", [python, "-m", "pytest", *pytest_args, "-q"], env)


if __name__ == "__main__":
    main()