MoviePilot/tests/run.py

"""
后端单测入口：默认并行执行文件分片，也支持单分片和显式串行调试。
"""
from __future__ import annotations

import argparse
import json
import subprocess
import sys
from pathlib import Path
from typing import Mapping, Sequence

import pytest

TESTS_DIR = Path(__file__).resolve().parent
RUNNER_PATH = Path(__file__).resolve()
# 默认启动 4 个独立测试进程，每个进程使用自己的隔离配置目录。
DEFAULT_SHARD_COUNT = 4
DURATIONS_PATH = TESTS_DIR / "fixtures" / "durations.json"
# 缺少历史耗时的文件按 1 秒参与分片负载估算。
DEFAULT_TEST_DURATION = 1.0
# CI 架构门禁任务单独运行的测试文件；分片使用 --exclude-architecture-gate 时不再重复执行
ARCHITECTURE_GATE_TESTS = (
    "test_architecture_dependencies.py",
    "test_architecture_adapter_imports.py",
    "test_architecture_egress.py",
    "test_architecture_event_facts.py",
    "test_architecture_event_policy.py",
    "test_architecture_contract_baseline.py",
    "test_architecture_baseline_cli.py",
)


def collect_test_files(exclude_architecture_gate: bool = False) -> list[Path]:
    """
    按稳定路径顺序收集测试根目录中的测试文件。

    :param exclude_architecture_gate: 是否排除已由 CI 架构门禁单独运行的文件
    :return: 排序后的 test_*.py 路径；不递归子目录
    """
    excluded = set(ARCHITECTURE_GATE_TESTS) if exclude_architecture_gate else set()
    return sorted(path for path in TESTS_DIR.glob("test_*.py") if path.name not in excluded)


def load_test_durations() -> dict[str, float]:
    """
    读取版本控制中的测试文件耗时估算。

    不使用本机运行缓存，确保本地与 CI 对相同版本作出相同分片决策。

    :return: 测试文件名到预计秒数的映射
    """
    return json.loads(DURATIONS_PATH.read_text(encoding="utf-8"))["files"]


def split_test_files(
    test_files: Sequence[Path], shard_count: int,
    durations: Mapping[str, float] | None = None,
) -> list[list[Path]]:
    """
    按耗时从长到短分配文件，优先填入预计最短的分片。

    未记录耗时的文件按 1 秒估算；耗时相同时按路径和较小分片编号稳定选择。

    :raises ValueError: 分片数小于等于 0
    :param test_files: 待分片的测试文件路径
    :param shard_count: 目标分片数，必须大于 0
    :param durations: 文件名到预计秒数的映射；None 时读取仓库中的耗时数据
    :return: 按分片编号排列的文件列表，每个分片内部再按路径排序
    """
    if shard_count <= 0:
        raise ValueError("shard_count 必须大于 0")
    durations = load_test_durations() if durations is None else durations
    shards: list[list[Path]] = [[] for _ in range(shard_count)]
    shard_durations = [0.0] * shard_count
    weighted_files = sorted(
        test_files,
        key=lambda path: (-durations.get(path.name, DEFAULT_TEST_DURATION), path),
    )
    for test_file in weighted_files:
        shard_index = min(range(shard_count), key=shard_durations.__getitem__)
        shards[shard_index].append(test_file)
        shard_durations[shard_index] += durations.get(test_file.name, DEFAULT_TEST_DURATION)
    return [sorted(shard) for shard in shards]


def parse_shard(value: str) -> tuple[int, int]:
    """
    解析从 1 开始编号的 N/TOTAL 分片标识。

    :raises argparse.ArgumentTypeError: 格式无效或不满足 1 <= N <= TOTAL
    :param value: 分片编号及总数文本，例如 2/4
    :return: 分片编号与总数
    """
    try:
        index_text, count_text = value.split("/", maxsplit=1)
        index = int(index_text)
        count = int(count_text)
    except (TypeError, ValueError) as error:
        raise argparse.ArgumentTypeError("--shard 必须使用 N/TOTAL 格式") from error
    if count <= 0 or not 1 <= index <= count:
        raise argparse.ArgumentTypeError("--shard 必须满足 1 <= N <= TOTAL")
    return index, count


def parse_args(argv: Sequence[str]) -> tuple[argparse.Namespace, list[str]]:
    """
    解析测试入口选项，将未识别选项留给 pytest。

    串行模式不能同时指定分片或排除架构门禁；参数错误由 argparse 终止并报告。

    :param argv: 不含脚本名称的命令行参数
    :return: 入口选项对象和待透传的 pytest 参数列表
    """
    parser = argparse.ArgumentParser(
        description="默认以 4 个独立 pytest 文件分片并行运行后端全量测试。"
    )
    parser.add_argument(
        "--serial",
        action="store_true",
        help="在当前进程串行运行 tests 目录，适合断点和顺序污染调试。",
    )
    parser.add_argument(
        "--shard",
        type=parse_shard,
        metavar="N/TOTAL",
        help="只运行指定文件分片；CI 使用同一参数启动独立 job。",
    )
    parser.add_argument(
        "--exclude-architecture-gate",
        action="store_true",
        help="跳过 CI 架构门禁已运行的测试文件，供覆盖率分片使用。",
    )
    args, pytest_args = parser.parse_known_args(argv)
    if args.serial and args.shard is not None:
        parser.error("--serial 不能与 --shard 同时使用")
    if args.serial and args.exclude_architecture_gate:
        parser.error("--serial 不能与 --exclude-architecture-gate 同时使用")
    return args, pytest_args


def run_pytest(paths: Sequence[Path], pytest_args: Sequence[str]) -> int:
    """
    在当前进程调用 pytest 执行指定路径。

    :param paths: 完整测试目录或选定分片的测试文件
    :param pytest_args: 原样附加的 pytest 参数
    :return: pytest 返回的退出状态码
    """
    return pytest.main([*(str(path) for path in paths), *pytest_args])


def _worker_command(
    shard_index: int, shard_count: int, pytest_args: Sequence[str],
    exclude_architecture_gate: bool = False,
) -> list[str]:
    """
    生成与 CI 使用相同分片算法的子进程命令。

    子进程按编号重新计算文件分片，排除选项必须与父进程一致。

    :param shard_index: 从 1 开始的分片编号
    :param shard_count: 总分片数
    :param pytest_args: 透传给子进程 pytest 的参数
    :param exclude_architecture_gate: 是否在子进程中同样排除架构门禁文件
    :return: 使用当前 Python 解释器的命令参数列表
    """
    return [
        sys.executable,
        str(RUNNER_PATH),
        "--shard",
        f"{shard_index}/{shard_count}",
        *(["--exclude-architecture-gate"] if exclude_architecture_gate else []),
        *pytest_args,
    ]


def run_parallel_shards(
    shards: Sequence[Sequence[Path]], pytest_args: Sequence[str],
    exclude_architecture_gate: bool = False,
) -> int:
    """
    为非空分片启动独立 pytest 进程并等待全部结束。

    收到键盘中断后终止仍在运行的子进程，并等待它们退出。

    :param shards: 用于启动和展示的分片列表；子进程依同一算法重新计算文件清单
    :param pytest_args: 传递给每个子进程的 pytest 参数
    :param exclude_architecture_gate: 父子进程共同使用的架构测试排除选项
    :return: 全部成功时为 0，否则返回按分片顺序遇到的首个非零退出码；中断时为 130
    """
    shard_count = len(shards)
    processes: list[tuple[int, subprocess.Popen]] = []
    for shard_index, shard in enumerate(shards, start=1):
        if not shard:
            continue
        print(
            f"启动测试分片 {shard_index}/{shard_count}：{len(shard)} 个文件",
            flush=True,
        )
        processes.append((
            shard_index,
            subprocess.Popen(_worker_command(
                shard_index, shard_count, pytest_args, exclude_architecture_gate,
            )),
        ))

    exit_code = 0
    try:
        for shard_index, process in processes:
            return_code = process.wait()
            if return_code != 0:
                print(
                    f"测试分片 {shard_index}/{shard_count} 失败，退出码 {return_code}",
                    file=sys.stderr,
                    flush=True,
                )
                exit_code = exit_code or return_code
    except KeyboardInterrupt:
        for _, process in processes:
            if process.poll() is None:
                process.terminate()
        for _, process in processes:
            process.wait()
        return 130
    return exit_code


def main(argv: Sequence[str] | None = None) -> int:
    """
    选择串行全量、指定分片或默认四分片并行执行。

    :param argv: 入口参数；None 时读取 sys.argv，空列表表示不传额外参数
    :return: 执行状态码；没有测试文件或所选分片为空时为 2
    """
    args, pytest_args = parse_args(sys.argv[1:] if argv is None else argv)
    if args.serial:
        return run_pytest([TESTS_DIR], pytest_args)

    test_files = collect_test_files(args.exclude_architecture_gate)
    if not test_files:
        print(f"未在 {TESTS_DIR} 找到 test_*.py", file=sys.stderr)
        return 2

    if args.shard is not None:
        shard_index, shard_count = args.shard
        selected = split_test_files(test_files, shard_count)[shard_index - 1]
        if not selected:
            print(f"测试分片 {shard_index}/{shard_count} 为空", file=sys.stderr)
            return 2
        print(
            f"运行测试分片 {shard_index}/{shard_count}：{len(selected)} 个文件",
            flush=True,
        )
        return run_pytest(selected, pytest_args)

    shards = split_test_files(test_files, DEFAULT_SHARD_COUNT)
    return run_parallel_shards(shards, pytest_args, args.exclude_architecture_gate)


if __name__ == "__main__":
    sys.exit(main())