# -*- coding: utf-8 -*- """ 灵活的测试执行脚本,可像搭积木一样组合不同参数或预置命令(模式/数据库/归档路径等), 直接运行本文件即可触发 pytest。 示例: python scripts/run_tests.py --suite online --mode ONLINE --keyword ORDERS python scripts/run_tests.py --preset offline_realdb python scripts/run_tests.py --suite online offline --db-dsn ... --json-archive 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__), "..")) SUITE_MAP: Dict[str, str] = { "online": "tests/unit/test_etl_tasks_online.py", "offline": "tests/unit/test_etl_tasks_offline.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( "--mode", choices=["ONLINE", "OFFLINE"], help="覆盖 TEST_MODE(默认沿用 .env / 环境变量)", ) parser.add_argument("--db-dsn", help="设置 TEST_DB_DSN,连接真实数据库进行测试") parser.add_argument("--json-archive", help="设置 TEST_JSON_ARCHIVE_DIR(离线档案目录)") parser.add_argument("--json-temp", help="设置 TEST_JSON_TEMP_DIR(临时 JSON 路径)") 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.mode: env_updates["TEST_MODE"] = args.mode if args.db_dsn: env_updates["TEST_DB_DSN"] = args.db_dsn if args.json_archive: env_updates["TEST_JSON_ARCHIVE_DIR"] = args.json_archive if args.json_temp: env_updates["TEST_JSON_TEMP_DIR"] = args.json_temp 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: # 默认跑 online + offline 套件 targets = [SUITE_MAP["online"], SUITE_MAP["offline"]] 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())