live-hub-py/tests/support/db.py

"""测试专用的 MySQL 数据库引擎与会话工具。"""

from __future__ import annotations

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

from sqlalchemy import Engine, create_engine, inspect, text
from sqlalchemy.exc import OperationalError
from sqlalchemy.orm import Session, sessionmaker
from sqlalchemy.schema import CreateColumn


def get_mysql_credentials() -> dict[str, str]:
    """获取连接 MySQL 所需的凭据（优先环境变量，其次项目 .env）。"""
    root_dir = Path(__file__).resolve().parents[2]
    env_file = root_dir / ".env"
    env_vars: dict[str, str] = {}
    if env_file.exists():
        with open(env_file, encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith("#") and "=" in line:
                    k, v = line.split("=", 1)
                    env_vars[k.strip()] = v.strip()

    user = os.getenv("MYSQL_USER") or env_vars.get("MYSQL_USER") or "douyu_login"
    pwd = os.getenv("MYSQL_PASSWORD") or env_vars.get("MYSQL_PASSWORD") or ""
    host = (
        os.getenv("MYSQL_BIND_HOST")
        or env_vars.get("MYSQL_BIND_HOST")
        or os.getenv("DEV_DB_HOST")
        or "127.0.0.1"
    )
    port = (
        os.getenv("MYSQL_HOST_PORT")
        or env_vars.get("MYSQL_HOST_PORT")
        or os.getenv("DEV_DB_PORT")
        or "3307"
    )
    db_name = (
        os.getenv("TEST_MYSQL_DATABASE")
        or env_vars.get("TEST_MYSQL_DATABASE")
        or "douyu_login_test"
    )
    assert_is_test_database(db_name)
    return {
        "user": user,
        "password": pwd,
        "host": host,
        "port": port,
        "database": db_name,
    }


def assert_is_test_database(db_name: str) -> None:
    """拒绝把测试指向开发库。

    ``clear_tables()`` 会 DELETE 掉所有业务表；一旦 TEST_MYSQL_DATABASE 被指到
    ``douyu_login``（开发库）甚至生产库，就是数据灾难。这里做一道硬拦截：
    库名与 .env 里的 MYSQL_DATABASE 相同就直接拒绝。
    """
    root_dir = Path(__file__).resolve().parents[2]
    dev_db = ""
    env_file = root_dir / ".env"
    if env_file.exists():
        with open(env_file, encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line.startswith("MYSQL_DATABASE="):
                    dev_db = line.split("=", 1)[1].strip()
                    break
    dev_db = os.getenv("MYSQL_DATABASE") or dev_db
    if dev_db and db_name == dev_db:
        raise RuntimeError(
            f"测试库名 ({db_name}) 与开发库同名，拒绝运行："
            "测试会清空所有表。请改用独立的测试库（默认 douyu_login_test）。"
        )


def get_test_database_url() -> str:
    """返回测试数据库的完整 URL。"""
    creds = get_mysql_credentials()
    return (
        f"mysql+pymysql://{quote_plus(creds['user'])}:{quote_plus(creds['password'])}"
        f"@{creds['host']}:{creds['port']}/{creds['database']}?charset=utf8mb4"
    )


# 1. 确保环境变量始终指向测试库
TEST_DB_URL = get_test_database_url()
os.environ["DATABASE_URL"] = TEST_DB_URL
os.environ["DB_NAME"] = get_mysql_credentials()["database"]

import web.backend.database as database_module
from web.backend.database import Base

# 2. 覆盖 database_module 中的变量，保证即使被提前 import 过也能指向 MySQL 测试库
database_module.DATABASE_URL = TEST_DB_URL

_test_engine: Engine | None = None


def get_test_engine() -> Engine:
    """获取测试库共享单例 Engine。"""
    global _test_engine
    if _test_engine is None:
        _test_engine = create_engine(
            TEST_DB_URL,
            pool_pre_ping=True,
            pool_size=10,
            max_overflow=20,
            pool_recycle=1800,
        )
        database_module.engine = _test_engine
        database_module.SessionLocal.configure(bind=_test_engine)
    return _test_engine


def make_test_session(engine: Engine | None = None) -> Session:
    """为测试创建一个独立的 Session。"""
    eng = engine or get_test_engine()
    return sessionmaker(bind=eng)()


def clear_tables(engine: Engine | None = None) -> None:
    """快速清空当前 Base.metadata 中的所有表数据（耗时约 0.01s）。"""
    eng = engine or get_test_engine()
    tables = [t.name for t in Base.metadata.sorted_tables]
    with eng.begin() as conn:
        conn.execute(text("SET FOREIGN_KEY_CHECKS = 0;"))
        for table in reversed(tables):
            conn.execute(text(f"DELETE FROM `{table}`;"))
        conn.execute(text("SET FOREIGN_KEY_CHECKS = 1;"))


def _add_missing_columns(engine: Engine) -> None:
    """给“已存在但结构偏旧”的测试表补列。

    ``Base.metadata.create_all`` 只建缺失的表，不会给已有表加列；测试库复用旧结构时，
    新模型字段会在 INSERT 阶段报 Unknown column。这里只做加法（不动索引/数据），
    让测试库跟模型对齐，又不会像 drop_all 那样丢掉迁移里才有的索引。
    """
    inspector = inspect(engine)
    with engine.begin() as conn:
        for table in Base.metadata.sorted_tables:
            if not inspector.has_table(table.name):
                continue
            existing = {column["name"] for column in inspector.get_columns(table.name)}
            for column in table.columns:
                if column.name in existing:
                    continue
                ddl = CreateColumn(column).compile(dialect=engine.dialect)
                conn.execute(text(f"ALTER TABLE `{table.name}` ADD COLUMN {ddl}"))


def init_test_schema(engine: Engine | None = None) -> None:
    """确保测试库中所有表结构存在并清空所有数据（幂等高效）。"""
    eng = engine or get_test_engine()
    try:
        _add_missing_columns(eng)
        Base.metadata.create_all(eng)
        clear_tables(eng)
    except OperationalError as exc:
        # 测试不再使用 SQLite，连不上库时必须给出可操作的提示，
        # 而不是抛一段 SQLAlchemy 原始堆栈。
        raise RuntimeError(
            "连接测试用 MySQL 失败，测试无法运行。\n"
            f"  连接串: {TEST_DB_URL}\n"
            "  请先启动数据库：./dev.sh（会起 docker compose，MySQL 映射到 127.0.0.1:3307）\n"
            "  或用 TEST_MYSQL_DATABASE / MYSQL_USER / MYSQL_PASSWORD / MYSQL_HOST_PORT "
            "指向一个可用的 MySQL 测试库。"
        ) from exc