xiaoxia-saas/tests/integration/conftest.py

"""
集成测试公共 fixtures

提供性能测试相关的工具、fixture 和 marker。
支持 pytest-xdist 并行执行：每个 worker 使用独立数据库，数据完全隔离。
"""

from __future__ import annotations

import os
import sys
import time
from contextlib import contextmanager
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, Dict, List, Optional

import pytest

# ── xdist 并行数据库隔离 ──────────────────────────────────────────────────
# 每个 xdist worker 进程创建独立的数据库并执行迁移，确保测试数据完全隔离
# 通过 PYTEST_XDIST_WORKER 环境变量识别 worker（如 gw0, gw1, ...）

_WORKER_DB_NAME: Optional[str] = None


def _get_worker_id() -> Optional[str]:
    """获取当前 xdist worker ID，非 worker 模式返回 None"""
    return os.environ.get("PYTEST_XDIST_WORKER")


def _parse_database_url(url: str) -> Dict[str, str]:
    """
    解析 DATABASE_URL，返回各组件。
    支持 postgresql+psycopg://user:pass@host:port/dbname 格式
    """
    from urllib.parse import urlparse

    parsed = urlparse(url)
    return {
        "driver": parsed.scheme,
        "user": parsed.username or "",
        "password": parsed.password or "",
        "host": parsed.hostname or "",
        "port": str(parsed.port or 5432),
        "dbname": parsed.path.lstrip("/") or "",
    }


def _create_worker_database(worker_id: str) -> str:
    """
    为 xdist worker 创建独立数据库并执行迁移。
    返回新的 DATABASE_URL。
    """
    base_url = os.environ.get(
        "DATABASE_URL",
        "postgresql+psycopg://postgres:postgres@localhost:5432/xiaoxia_saas",
    )
    db_info = _parse_database_url(base_url)

    # 生成 worker 专属数据库名
    base_db = db_info["dbname"]
    worker_db = f"{base_db}_{worker_id}"
    global _WORKER_DB_NAME
    _WORKER_DB_NAME = worker_db

    # 使用 psycopg 创建数据库（连接到 postgres 库）
    try:
        import psycopg

        conn_str = (
            f"host={db_info['host']} port={db_info['port']} "
            f"user={db_info['user']} password={db_info['password']} "
            f"dbname=postgres"
        )
        conn = psycopg.connect(conn_str, autocommit=True)
        cur = conn.cursor()

        # 先尝试删除（防止残留）
        cur.execute(f'DROP DATABASE IF EXISTS "{worker_db}" WITH (FORCE)')

        # 创建新数据库
        cur.execute(f'CREATE DATABASE "{worker_db}"')
        cur.close()
        conn.close()
        print(f"[xdist {worker_id}] ✅ 创建数据库: {worker_db}")
    except ImportError:
        print(f"[xdist {worker_id}] ⚠️  psycopg 未安装，跳过数据库创建")
        return base_url
    except Exception as e:
        print(f"[xdist {worker_id}] ⚠️  创建数据库失败: {e}")
        return base_url

    # 构建新的 DATABASE_URL
    new_url = (
        f"{db_info['driver']}://{db_info['user']}:{db_info['password']}"
        f"@{db_info['host']}:{db_info['port']}/{worker_db}"
    )

    # 执行 alembic 迁移
    print(f"[xdist {worker_id}] 🔄 执行 Alembic 迁移...")
    try:
        ROOT = Path(__file__).resolve().parents[2]
        api_path = str(ROOT / "apps" / "api")
        if api_path not in sys.path:
            sys.path.insert(0, api_path)
        if str(ROOT) not in sys.path:
            sys.path.insert(0, str(ROOT))

        from alembic import command as alembic_command
        from alembic.config import Config as AlembicConfig

        alembic_cfg = AlembicConfig(str(ROOT / "alembic.ini"))
        alembic_cfg.set_main_option("sqlalchemy.url", new_url)
        # 兼容不同的脚本路径配置
        alembic_cfg.set_main_option("script_location", str(ROOT / "alembic"))

        # 临时设置环境变量供 alembic env.py 使用
        os.environ["DATABASE_URL"] = new_url
        alembic_command.upgrade(alembic_cfg, "head")
        print(f"[xdist {worker_id}] ✅ 迁移完成")
    except Exception as e:
        print(f"[xdist {worker_id}] ❌ 迁移失败: {e}")
        raise

    return new_url


def _cleanup_worker_database(worker_id: str):
    """清理 xdist worker 的数据库"""
    global _WORKER_DB_NAME
    if not _WORKER_DB_NAME:
        return

    base_url = os.environ.get(
        "DATABASE_URL",
        "postgresql+psycopg://postgres:postgres@localhost:5432/xiaoxia_saas",
    )
    db_info = _parse_database_url(base_url)

    try:
        import psycopg

        conn_str = (
            f"host={db_info['host']} port={db_info['port']} "
            f"user={db_info['user']} password={db_info['password']} "
            f"dbname=postgres"
        )
        conn = psycopg.connect(conn_str, autocommit=True)
        cur = conn.cursor()
        # 强制断开所有连接后删除
        cur.execute(
            f"SELECT pg_terminate_backend(pid) FROM pg_stat_activity "
            f"WHERE datname = '{_WORKER_DB_NAME}' AND pid <> pg_backend_pid()"
        )
        cur.execute(f'DROP DATABASE IF EXISTS "{_WORKER_DB_NAME}" WITH (FORCE)')
        cur.close()
        conn.close()
        print(f"[xdist {worker_id}] 🧹 已清理数据库: {_WORKER_DB_NAME}")
    except Exception as e:
        print(f"[xdist {worker_id}] ⚠️  清理数据库失败: {e}")
    finally:
        _WORKER_DB_NAME = None


# ── 性能阈值配置 ──────────────────────────────────────────────────────────
PERF_THRESHOLDS: Dict[str, int] = {
    "core": 500,  # 核心接口：500ms
    "normal": 1000,  # 普通接口：1000ms
    "heavy": 3000,  # 重操作：3000ms（涉及外部调用或复杂计算）
}

# 性能测试是否跳过（通过环境变量控制）
SKIP_PERF_TESTS = os.environ.get("SKIP_PERF_TESTS", "").lower() in ("1", "true", "yes")

# 性能测试容忍度：允许一定比例的请求超标（避免CI偶发波动）
# 默认：3次请求中允许1次超标（取中位数判断）
PERF_SAMPLE_COUNT = int(os.environ.get("PERF_SAMPLE_COUNT", "3"))
PERF_TOLERANCE_RATIO = float(os.environ.get("PERF_TOLERANCE_RATIO", "0.34"))


# ── 数据类 ────────────────────────────────────────────────────────────────


@dataclass
class PerfResult:
    """单次性能测试结果"""

    name: str
    threshold_ms: int
    times_ms: List[float] = field(default_factory=list)
    status_code: Optional[int] = None

    @property
    def median_ms(self) -> float:
        if not self.times_ms:
            return 0.0
        sorted_times = sorted(self.times_ms)
        n = len(sorted_times)
        if n % 2 == 0:
            return (sorted_times[n // 2 - 1] + sorted_times[n // 2]) / 2
        return sorted_times[n // 2]

    @property
    def mean_ms(self) -> float:
        if not self.times_ms:
            return 0.0
        return sum(self.times_ms) / len(self.times_ms)

    @property
    def min_ms(self) -> float:
        return min(self.times_ms) if self.times_ms else 0.0

    @property
    def max_ms(self) -> float:
        return max(self.times_ms) if self.times_ms else 0.0

    @property
    def passed(self) -> bool:
        """判断是否通过：基于中位数 + 容忍比例"""
        if not self.times_ms:
            return False
        # 中位数必须在阈值内
        if self.median_ms > self.threshold_ms:
            return False
        # 超标比例不能超过容忍度
        over_count = sum(1 for t in self.times_ms if t > self.threshold_ms)
        over_ratio = over_count / len(self.times_ms)
        return over_ratio <= PERF_TOLERANCE_RATIO


# ── 性能断言上下文管理器 ──────────────────────────────────────────────────


class PerfAssert:
    """
    性能断言工具。

    使用方式：
        def test_login_performance(client, perf_assert):
            with perf_assert("core", name="login") as result:
                response = client.post("/api/v1/auth/login", json={...})
                result.status_code = response.status_code
            # 退出 with 块时自动断言
    """

    def __init__(self, sample_count: int = PERF_SAMPLE_COUNT):
        self.sample_count = sample_count
        self.results: List[PerfResult] = []

    @contextmanager
    def __call__(self, threshold_level: str = "core", name: str = "", samples: Optional[int] = None):
        """
        创建一个性能测试上下文。

        Args:
            threshold_level: 阈值级别 ("core", "normal", "heavy")
            name: 测试名称（用于输出报告）
            samples: 采样次数，默认使用全局配置
        """
        if threshold_level not in PERF_THRESHOLDS:
            raise ValueError(f"未知的阈值级别: {threshold_level}，可选: {list(PERF_THRESHOLDS.keys())}")

        threshold_ms = PERF_THRESHOLDS[threshold_level]
        result = PerfResult(name=name or threshold_level, threshold_ms=threshold_ms)

        # 预热（第一次请求可能有冷启动开销）
        yield result
        # 第一次调用已经记录在 result.times_ms 中（由调用方通过 measure 方法）

    def measure(self, threshold_level: str = "core", name: str = "", samples: Optional[int] = None) -> Callable:
        """
        返回一个装饰器/包装器，用于测量函数执行时间。

        使用方式：
            result = perf_assert.measure("core", "login")(
                lambda: client.post("/api/v1/auth/login", json={...})
            )
        """

        def wrapper(func):
            if threshold_level not in PERF_THRESHOLDS:
                raise ValueError(f"未知的阈值级别: {threshold_level}，可选: {list(PERF_THRESHOLDS.keys())}")
            threshold_ms = PERF_THRESHOLDS[threshold_level]
            num_samples = samples or self.sample_count
            result = PerfResult(name=name or threshold_level, threshold_ms=threshold_ms)

            last_response = None
            for _ in range(num_samples):
                start = time.perf_counter()
                last_response = func()
                elapsed = (time.perf_counter() - start) * 1000
                result.times_ms.append(elapsed)

            if hasattr(last_response, "status_code"):
                result.status_code = last_response.status_code

            self.results.append(result)
            return result

        return wrapper

    def assert_all(self):
        """断言所有性能测试结果都通过"""
        failed = [r for r in self.results if not r.passed]
        if failed:
            lines = []
            for r in failed:
                lines.append(
                    f"  ❌ {r.name}: 中位数 {r.median_ms:.1f}ms "
                    f"(阈值 {r.threshold_ms}ms) "
                    f"[min={r.min_ms:.1f}, max={r.max_ms:.1f}, "
                    f"mean={r.mean_ms:.1f}, samples={len(r.times_ms)}]"
                )
            raise AssertionError(f"性能测试失败 ({len(failed)}/{len(self.results)}):\n" + "\n".join(lines))

    def report(self) -> str:
        """生成性能报告文本"""
        lines = ["=" * 60, " 性能测试报告", "=" * 60]
        for r in self.results:
            status = "✅" if r.passed else "❌"
            lines.append(f" {status} {r.name:<40s} " f"median={r.median_ms:>7.1f}ms / {r.threshold_ms:>5d}ms")
            lines.append(
                f"     min={r.min_ms:.1f}ms  max={r.max_ms:.1f}ms  "
                f"mean={r.mean_ms:.1f}ms  samples={len(r.times_ms)}"
                f"  status={r.status_code or 'N/A'}"
            )
        passed = sum(1 for r in self.results if r.passed)
        lines.append("=" * 60)
        lines.append(f" 总计: {passed}/{len(self.results)} 通过")
        lines.append("=" * 60)
        return "\n".join(lines)


# ── pytest hooks ──────────────────────────────────────────────────────────


def pytest_configure(config):
    """
    pytest 配置钩子。

    - 注册自定义 marker
    - xdist worker 模式下：创建独立数据库 + 执行迁移
    """
    # 注册自定义 marker
    config.addinivalue_line("markers", "performance: 标记为性能测试（可通过 -m 'not performance' 跳过）")
    config.addinivalue_line("markers", "perf_core: 核心接口性能测试（阈值 500ms）")
    config.addinivalue_line("markers", "perf_normal: 普通接口性能测试（阈值 1000ms）")
    config.addinivalue_line("markers", "perf_heavy: 重操作接口性能测试（阈值 3000ms）")

    # xdist worker 模式：创建独立数据库并执行迁移
    worker_id = _get_worker_id()
    if worker_id:
        # 只有当 USE_IN_MEMORY_DB 不为 true 时才创建独立数据库
        use_in_memory = os.environ.get("USE_IN_MEMORY_DB", "true").lower() == "true"
        if not use_in_memory:
            print(f"[xdist {worker_id}] 🚀 worker 启动，准备独立数据库...")
            new_db_url = _create_worker_database(worker_id)
            os.environ["DATABASE_URL"] = new_db_url
        else:
            print(f"[xdist {worker_id}] ℹ️  USE_IN_MEMORY_DB=true，跳过 worker 数据库创建")


def pytest_unconfigure(config):
    """pytest 结束钩子：清理 xdist worker 数据库"""
    worker_id = _get_worker_id()
    if worker_id and _WORKER_DB_NAME:
        _cleanup_worker_database(worker_id)


def pytest_collection_modifyitems(config, items):
    """根据环境变量自动跳过性能测试"""
    if SKIP_PERF_TESTS:
        skip_perf = pytest.mark.skip(reason="SKIP_PERF_TESTS=1，跳过性能测试")
        for item in items:
            if "performance" in item.keywords or "perf_" in item.keywords:
                item.add_marker(skip_perf)


# ── pytest fixtures ──────────────────────────────────────────────────────


@pytest.fixture
def perf_assert():
    """
    性能断言 fixture。

    使用方式 1（推荐，自动断言）：
        def test_login(client, perf_assert):
            @perf_assert.measure("core", "POST /auth/login")
            def _call():
                return client.post("/api/v1/auth/login", json={...})

            result = _call()
            assert result.status_code == 200

    使用方式 2（手动多次调用）：
        def test_login(client, perf_assert):
            result = perf_assert.run("core", "POST /auth/login",
                lambda: client.post("/api/v1/auth/login", json={...})
            )
            assert result.status_code == 200
    """
    return PerfAssert()


@pytest.fixture
def perf_thresholds():
    """返回性能阈值配置字典"""
    return dict(PERF_THRESHOLDS)


# ── 辅助函数 ──────────────────────────────────────────────────────────────


def run_perf_test(
    name: str,
    threshold_level: str,
    func: Callable,
    samples: int = PERF_SAMPLE_COUNT,
) -> PerfResult:
    """
    运行一次性能测试（独立函数，方便在 fixture 外部使用）。

    Args:
        name: 测试名称
        threshold_level: 阈值级别
        func: 要测量的函数（无参数）
        samples: 采样次数

    Returns:
        PerfResult 对象
    """
    if threshold_level not in PERF_THRESHOLDS:
        raise ValueError(f"未知的阈值级别: {threshold_level}，可选: {list(PERF_THRESHOLDS.keys())}")

    threshold_ms = PERF_THRESHOLDS[threshold_level]
    result = PerfResult(name=name, threshold_ms=threshold_ms)

    last_response = None
    for _ in range(samples):
        start = time.perf_counter()
        last_response = func()
        elapsed = (time.perf_counter() - start) * 1000
        result.times_ms.append(elapsed)

    if hasattr(last_response, "status_code"):
        result.status_code = last_response.status_code

    return result