feiqiu-ETL/etl_billiards/scripts/run_tests.py

# -*- coding: utf-8 -*-
"""
灵活的测试执行脚本，可像搭积木一样组合不同参数或预置命令（模式/数据库/归档路径等），
直接运行本文件即可触发 pytest。

示例：
    python scripts/run_tests.py --suite online --flow FULL --keyword ORDERS
    python scripts/run_tests.py --preset fetch_only
    python scripts/run_tests.py --suite online --json-source tmp/archives
"""
from __future__ import annotations

import argparse
import importlib.util
import os
import shlex
import sys
from typing import Dict, List

import pytest

PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))

# 确保项目根目录在 sys.path，便于 tests 内部 import config / tasks 等模块
if PROJECT_ROOT not in sys.path:
    sys.path.insert(0, PROJECT_ROOT)

SUITE_MAP: Dict[str, str] = {
    "online": "tests/unit/test_etl_tasks_online.py",
    "integration": "tests/integration/test_database.py",
}

PRESETS: Dict[str, Dict] = {}


def _load_presets():
    preset_path = os.path.join(os.path.dirname(__file__), "test_presets.py")
    if not os.path.exists(preset_path):
        return
    spec = importlib.util.spec_from_file_location("test_presets", preset_path)
    if not spec or not spec.loader:
        return
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)  # type: ignore[attr-defined]
    presets = getattr(module, "PRESETS", {})
    if isinstance(presets, dict):
        PRESETS.update(presets)


_load_presets()

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="ETL 测试执行器（支持参数化调配）")
    parser.add_argument(
        "--suite",
        choices=sorted(SUITE_MAP.keys()),
        nargs="+",
        help="预置测试套件，可多选（默认全部 online/offline）",
    )
    parser.add_argument(
        "--tests",
        nargs="+",
        help="自定义测试路径（可与 --suite 混用），例如 tests/unit/test_config.py",
    )
    parser.add_argument(
        "--flow",
        choices=["FETCH_ONLY", "INGEST_ONLY", "FULL"],
        help="覆盖 PIPELINE_FLOW（在线抓取/本地清洗/全流程）",
    )
    parser.add_argument("--json-source", help="设置 JSON_SOURCE_DIR（本地清洗入库使用的 JSON 目录）")
    parser.add_argument("--json-fetch-root", help="设置 JSON_FETCH_ROOT（在线抓取输出根目录）")
    parser.add_argument(
        "--keyword",
        "-k",
        help="pytest -k 关键字过滤（例如 ORDERS，只运行包含该字符串的用例）",
    )
    parser.add_argument(
        "--pytest-args",
        help="附加 pytest 参数，格式与命令行一致（例如 \"-vv --maxfail=1\"）",
    )
    parser.add_argument(
        "--env",
        action="append",
        metavar="KEY=VALUE",
        help="自定义环境变量，可重复传入，例如 --env STORE_ID=123",
    )
    parser.add_argument("--preset", choices=sorted(PRESETS.keys()) if PRESETS else None, nargs="+",
                        help="从 scripts/test_presets.py 中选择一个或多个组合命令")
    parser.add_argument("--list-presets", action="store_true", help="列出可用预置命令后退出")
    parser.add_argument("--dry-run", action="store_true", help="仅打印将要执行的命令与环境，不真正运行 pytest")
    return parser.parse_args()


def apply_presets_to_args(args: argparse.Namespace):
    if not args.preset:
        return
    for name in args.preset:
        preset = PRESETS.get(name, {})
        if not preset:
            continue
        for key, value in preset.items():
            if key in ("suite", "tests"):
                if not value:
                    continue
                existing = getattr(args, key)
                if existing is None:
                    setattr(args, key, list(value))
                else:
                    existing.extend(value)
            elif key == "env":
                args.env = (args.env or []) + list(value)
            elif key == "pytest_args":
                args.pytest_args = " ".join(filter(None, [value, args.pytest_args or ""]))
            elif key == "keyword":
                if args.keyword is None:
                    args.keyword = value
            else:
                if getattr(args, key, None) is None:
                    setattr(args, key, value)


def apply_env(args: argparse.Namespace) -> Dict[str, str]:
    env_updates = {}
    if args.flow:
        env_updates["PIPELINE_FLOW"] = args.flow
    if args.json_source:
        env_updates["JSON_SOURCE_DIR"] = args.json_source
    if args.json_fetch_root:
        env_updates["JSON_FETCH_ROOT"] = args.json_fetch_root
    if args.env:
        for item in args.env:
            if "=" not in item:
                raise SystemExit(f"--env 参数格式错误: {item!r}，应为 KEY=VALUE")
            key, value = item.split("=", 1)
            env_updates[key.strip()] = value.strip()

    for key, value in env_updates.items():
        os.environ[key] = value
    return env_updates


def build_pytest_args(args: argparse.Namespace) -> List[str]:
    targets: List[str] = []
    if args.suite:
        for suite in args.suite:
            targets.append(SUITE_MAP[suite])
    if args.tests:
        targets.extend(args.tests)
    if not targets:
        targets = list(SUITE_MAP.values())

    pytest_args: List[str] = targets
    if args.keyword:
        pytest_args += ["-k", args.keyword]
    if args.pytest_args:
        pytest_args += shlex.split(args.pytest_args)
    return pytest_args


def main() -> int:
    os.chdir(PROJECT_ROOT)
    args = parse_args()
    if args.list_presets:
        print("可用预置命令：")
        if not PRESETS:
            print("（暂无，可编辑 scripts/test_presets.py 添加）")
        else:
            for name in sorted(PRESETS):
                print(f"- {name}")
        return 0

    apply_presets_to_args(args)
    env_updates = apply_env(args)
    pytest_args = build_pytest_args(args)

    print("=== 环境变量覆盖 ===")
    if env_updates:
        for k, v in env_updates.items():
            print(f"{k}={v}")
    else:
        print("（无覆盖，沿用系统默认）")
    print("\n=== Pytest 参数 ===")
    print(" ".join(pytest_args))
    print()

    if args.dry_run:
        print("Dry-run 模式，未真正执行 pytest")
        return 0

    exit_code = pytest.main(pytest_args)
    return int(exit_code)


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