"""测试专用的 MySQL 数据库引擎与会话工具。""" from __future__ import annotations import os from pathlib import Path from urllib.parse import quote_plus from sqlalchemy import Engine, create_engine, text from sqlalchemy.exc import OperationalError from sqlalchemy.orm import Session, sessionmaker 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 init_test_schema(engine: Engine | None = None) -> None: """确保测试库中所有表结构存在并清空所有数据(幂等高效)。""" eng = engine or get_test_engine() try: 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