litellm/tests/integration/run.py

from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
from pathlib import Path
from types import MappingProxyType
from typing import Final

GROUPS: Final = MappingProxyType(
    {
        "management": ("management", "authorization", "configuration"),
        "accounting": ("pricing", "spend"),
        "database": ("database",),
        "providers": ("providers", "routing", "streaming"),
        "extensions": ("observability", "compatibility"),
        "mcp": ("mcp",),
        "sdk": ("sdk",),
        "cost": ("cost_calculation",),
    }
)


def main() -> int:
    parser: Final = argparse.ArgumentParser()
    parser.add_argument("group", choices=tuple(GROUPS))
    parser.add_argument("--results", type=Path, default=Path("test-results/integration"))
    parser.add_argument("--seed", type=int, default=int(os.environ.get("INTEGRATION_SEED", "4106601")))
    parser.add_argument("--order-seed", type=int, default=int(os.environ.get("INTEGRATION_ORDER_SEED", "0")))
    parser.add_argument("--workers", type=int, default=int(os.environ.get("INTEGRATION_WORKERS", "1")))
    options: Final = parser.parse_args()
    root: Final = Path(__file__).resolve().parents[2]
    selected: Final = tuple(
        str(path.relative_to(root))
        for folder in GROUPS[options.group]
        for path in sorted((root / "tests/integration" / folder).glob("test_*.py"))
    )
    if not selected:
        parser.error(f"No integration test files selected for {options.group}")
    output: Final = options.results.resolve()
    output.mkdir(parents=True, exist_ok=True)
    environment: Final = {
        **os.environ,
        "PYTHONPATH": os.pathsep.join((str(root), str(root / "tests"), str(root / "tests/e2e"))),
        "INTEGRATION_RESULTS_DIR": str(output),
        "LITELLM_LOCAL_MODEL_COST_MAP": "True",
    }
    result: Final = subprocess.call(
        [
            sys.executable,
            "-m",
            "pytest",
            *selected,
            "-vv",
            "-rs",
            "--strict-markers",
            "-p",
            "no:pytest-retry",
            "-p",
            "no:rerunfailures",
            "--timeout=90",
            "--durations=15",
            f"--hypothesis-seed={options.seed}",
            f"--integration-order-seed={options.order_seed}",
            f"--junitxml={output / 'junit.xml'}",
            *(("-n", str(options.workers)) if options.workers > 1 else ()),
        ],
        cwd=root,
        env=environment,
    )
    if result != 0:
        return result
    evidence: Final = json.loads((output / "execution.json").read_text())
    collected_files: Final = {node.split("::", 1)[0] for node in evidence["collected"]}
    empty: Final = tuple(path for path in selected if path not in collected_files)
    if empty:
        sys.stderr.write(f"Selected integration files collected zero tests: {', '.join(empty)}\n")
        return 1
    if not evidence["complete"]:
        sys.stderr.write("Integration run did not complete: a collected node neither passed nor skipped\n")
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())