xiaoxia-saas/scripts/ci/select_unit_tests.py

#!/usr/bin/env python3
"""
根据PR改动文件选择需要运行的单元测试文件。

映射规则：
1. 改了tests/unit下的测试文件 -> 直接跑这些测试
2. 改了apps/api/app/api/routes/xxx.py -> 匹配 test_*xxx*.py
3. 改了apps/api/app/services/xxx.py -> 匹配 test_*xxx*.py
4. 改了apps/worker/.../xxx.py -> 匹配 test_*xxx*.py
5. 改了apps/worker/video_processing/xxx_engine.py -> 匹配 test_*xxx*.py
6. 改了packages/.../xxx.py -> 匹配 test_*xxx*.py
7. 改了公共核心模块(core/middleware/schemas/config/db/auth/dependencies) -> 全量
8. 改了依赖文件(requirements*.txt, pyproject.toml) -> 全量
9. 改了alembic/migrations -> 全量
10. 匹配不到测试的改动 -> 全量兜底
"""

import os
import subprocess
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
TESTS_DIR = ROOT / "tests" / "unit"

# 触发全量的文件模式（公共核心/基础设施）
FULL_RUN_PATTERNS = [
    "apps/api/app/core/",
    "apps/api/app/middleware/",
    "apps/api/app/schemas/",
    "apps/api/app/config.py",
    "apps/api/app/db.py",
    "apps/api/app/auth.py",
    "apps/api/app/dependencies.py",
    "packages/shared/",
    "alembic/",
    "migrations/",
    "requirements-base.txt",
    "requirements.txt",
    "requirements-dev.txt",
    "pyproject.toml",
    "setup.cfg",
    ".gitea/workflows/",
    "scripts/ci/",
    "tests/conftest.py",
]

# 目录到测试文件关键词的映射（模糊匹配）
DIR_KEYWORD_MAP = {
    "apps/api/app/api/routes/": "",  # 用文件名匹配
    "apps/api/app/services/": "",  # 用文件名匹配
    "apps/worker/worker_app/tasks/": "",
    "apps/worker/video_processing/": "",
    "apps/worker/services/": "",
    "packages/application/": "",
    "packages/adapters/": "",
}


def get_changed_files():
    """获取改动文件列表（从环境变量或git diff）。"""
    # 优先从环境变量读取（CI中传入）
    changed_env = os.environ.get("CHANGED_FILES", "")
    if changed_env:
        return [f.strip() for f in changed_env.split("\n") if f.strip()]

    # 回退到git diff（本地调试用）
    try:
        result = subprocess.run(
            ["git", "diff", "--name-only", "origin/develop...HEAD"],
            capture_output=True,
            text=True,
            cwd=ROOT,
        )
        if result.returncode == 0:
            return [f.strip() for f in result.stdout.split("\n") if f.strip()]
    except Exception:
        pass

    return []


def should_full_run(files):
    """检查是否需要全量运行。"""
    for f in files:
        for pattern in FULL_RUN_PATTERNS:
            if f.startswith(pattern) or f == pattern:
                print(f"[full-run] 触发全量: {f} 匹配 {pattern}")
                return True
    return False


def extract_module_name(filepath):
    """从文件路径提取模块名（用于匹配测试文件）。"""
    # 去掉扩展名
    name = Path(filepath).stem

    # 特殊映射
    special_mappings = {
        # 路由文件
        "edit_plans_adjustments": "edit_plan_adjustments",
        "edit_plans_ai": "edit_plan",
        "edit_plans_clips_batch": "edit_plan",
        "edit_plans_cover": "edit_plan_cover",
        "edit_plans_export": "edit_plan_export",
        "edit_plans_filter": "edit_plan_filter",
        "edit_plans_generation": "edit_plan_generation",
        "edit_plans_transitions": "edit_plan_transitions",
        "asset_libraries": "asset_library",
        "classification_jobs": "classification",
        "chunked_upload": "chunked_upload",
        "form_upload": "form_upload",
        # 服务文件
        "edit_template_service": "edit_template_service",
        "edit_plan_service": "edit_plan_service",
        "unified_render_service": "unified_render",
        "job_service": "job_service",
        "auto_clip_service": "auto_clip",
        "cosyvoice_service": "cosyvoice",
        "video_compose_service": "video_compose",
        "email_service": "email_service",
    }

    return special_mappings.get(name, name)


def find_matching_tests(keyword, all_test_files):
    """模糊匹配测试文件。"""
    keyword_lower = keyword.lower().replace("_", "")
    matches = []
    for tf in all_test_files:
        tf_name = Path(tf).stem.lower().replace("_", "")
        if keyword_lower in tf_name or tf_name in keyword_lower:
            matches.append(tf)
    return matches


def get_all_test_files():
    """获取所有单元测试文件。"""
    if not TESTS_DIR.exists():
        return []
    return sorted(str(f.relative_to(ROOT)) for f in TESTS_DIR.glob("test_*.py"))


def select_tests(changed_files):
    """主函数：选择要运行的测试文件。"""
    all_tests = get_all_test_files()

    if not changed_files:
        print("[info] 未找到改动文件，全量运行")
        return all_tests, "full (no changes detected)"

    if should_full_run(changed_files):
        return all_tests, "full (core/common files changed)"

    selected = set()
    test_file_changes = []
    source_file_changes = []

    for f in changed_files:
        # 测试文件本身改动（仅保留仍存在的文件，删除的测试文件不加入运行列表）
        if f.startswith("tests/unit/test_") and f.endswith(".py"):
            if (ROOT / f).exists():
                test_file_changes.append(f)
                selected.add(f)
            else:
                print(f"[skip-deleted] 测试文件已删除，跳过: {f}")
        # 源码文件改动
        elif f.endswith(".py"):
            source_file_changes.append(f)
            module_name = extract_module_name(f)
            matches = find_matching_tests(module_name, all_tests)
            if matches:
                for m in matches:
                    selected.add(m)
                print(f"[map] {f} -> {len(matches)} 个测试: {[Path(m).name for m in matches]}")
            else:
                print(f"[nomatch] {f} (module: {module_name}) 未找到匹配的测试文件")

    if not selected:
        print("[info] 未匹配到任何测试文件，全量运行兜底")
        return all_tests, "full (no matching tests)"

    return sorted(selected), f"incremental ({len(selected)} test files)"


def main():
    changed_files = get_changed_files()
    print(f"=== 改动文件 ({len(changed_files)} 个) ===")
    for f in changed_files[:20]:
        print(f"  {f}")
    if len(changed_files) > 20:
        print(f"  ... 还有 {len(changed_files) - 20} 个")
    print()

    selected, mode = select_tests(changed_files)

    print()
    print(f"=== 运行模式: {mode} ===")
    print(f"=== 选中测试文件: {len(selected)} 个 ===")
    for t in selected[:20]:
        print(f"  {t}")
    if len(selected) > 20:
        print(f"  ... 还有 {len(selected) - 20} 个")

    # 输出结果文件（供CI后续步骤使用）
    output_file = os.environ.get("SELECTED_TESTS_OUTPUT", "")
    if output_file:
        with open(output_file, "w") as f:
            for t in selected:
                f.write(t + "\n")
        print(f"\n已写入到: {output_file}")

    # 设置环境变量标记
    gh_output = os.environ.get("GITHUB_OUTPUT", "")
    if gh_output:
        with open(gh_output, "a") as f:
            f.write(f"test_count={len(selected)}\n")
            f.write("mode=" + ("incremental" if "incremental" in mode else "full") + "\n")

    # 退出码：0=增量, 1=全量（供CI判断）
    sys.exit(0 if "incremental" in mode else 1)


if __name__ == "__main__":
    main()