litellm/tests/rust-python-harness/shared/unit_runners/python_runner.py

from __future__ import annotations

import argparse
import importlib
import os
import subprocess
import sys
import tempfile
from collections.abc import Callable, Sequence
from pathlib import Path
from typing import TYPE_CHECKING, Final, Literal, cast

from pluggy import HookimplMarker
from pydantic import BaseModel, ConfigDict

if TYPE_CHECKING:
    import pytest

hookimpl: Final = HookimplMarker("pytest")

Backend = Literal["python", "rust"]


class PythonReport(BaseModel):
    model_config = ConfigDict(frozen=True, extra="forbid")

    backend: Backend
    verified: bool
    tests: tuple[str, ...] = ()
    outcomes: tuple[tuple[str, str, str], ...] = ()
    exit_code: int
    problems: tuple[str, ...] = ()


class BackendSpec(BaseModel):
    model_config = ConfigDict(frozen=True, extra="forbid")

    environment_variable: str
    probe: str = ""


class WorkerArgs(BaseModel):
    model_config = ConfigDict(frozen=True, extra="forbid")

    backend: Backend
    probe: str
    output: Path
    pytest_args: tuple[str, ...]


class ResultPlugin:
    def __init__(self, backend: Backend, probe: Callable[[], object] | None) -> None:
        self.backend: Final = backend
        self.probe: Final = probe
        self.tests: tuple[str, ...] = ()
        self.outcomes: tuple[tuple[str, str, str], ...] = ()
        self.problems: tuple[str, ...] = ()

    def verify(self) -> None:
        if self.probe is not None and self.probe() != self.backend:
            raise RuntimeError(f"backend probe did not select {self.backend}")

    def pytest_collection_finish(self, session: pytest.Session) -> None:
        self.tests = tuple(item.nodeid for item in session.items)

    @hookimpl(tryfirst=True)
    def pytest_runtest_call(self, item: pytest.Item) -> None:
        del item
        self.verify()

    def pytest_collectreport(self, report: pytest.CollectReport) -> None:
        if report.failed:
            self.problems = (*self.problems, str(report.longrepr))

    def pytest_runtest_logreport(self, report: pytest.TestReport) -> None:
        self.outcomes = (*self.outcomes, (report.nodeid, report.when, report.outcome))
        if report.failed:
            self.problems = (*self.problems, str(report.longrepr))


def run_python_tests(
    selectors: Sequence[str],
    repo_root: Path,
    backend: Backend,
    spec: BackendSpec,
    pytest_args: Sequence[str] = (),
) -> PythonReport:
    with tempfile.TemporaryDirectory(prefix="litellm-unit-tests-") as directory:
        output: Final = Path(directory) / "report.json"
        command: Final = (
            sys.executable,
            "-m",
            __name__,
            "--backend",
            backend,
            *(("--probe", spec.probe) if spec.probe else ()),
            "--output",
            str(output),
            "--",
            *selectors,
            *pytest_args,
        )
        env: Final = {
            **os.environ,
            spec.environment_variable: "1" if backend == "rust" else "0",
            "PYTHONPATH": os.pathsep.join((str(repo_root), os.environ.get("PYTHONPATH", ""))),
        }
        try:
            result: Final = subprocess.run(
                command, cwd=repo_root, env=env, capture_output=True, text=True, timeout=600, check=False
            )
        except (OSError, subprocess.TimeoutExpired) as error:
            return PythonReport(backend=backend, verified=False, exit_code=1, problems=(str(error),))
        if not output.exists():
            return PythonReport(
                backend=backend,
                verified=False,
                exit_code=result.returncode or 1,
                problems=(result.stdout + result.stderr,),
            )
        report: Final = PythonReport.model_validate_json(output.read_text())
        process_output: Final = (result.stdout + result.stderr).strip()
        if result.returncode and not report.problems and process_output:
            return report.model_copy(update={"problems": (process_output,)})
        if report.exit_code != result.returncode:
            return report.model_copy(
                update={
                    "exit_code": result.returncode or 1,
                    "problems": (*report.problems, "worker exit code differs from report"),
                }
            )
        return report


def compare_python_runs(python: PythonReport, rust: PythonReport) -> tuple[str, ...]:
    python_only: Final = tuple(sorted(set(python.outcomes) - set(rust.outcomes)))
    rust_only: Final = tuple(sorted(set(rust.outcomes) - set(python.outcomes)))
    return (
        *(("backend selection was not verified",) if not python.verified or not rust.verified else ()),
        *(("Python run used the wrong backend",) if python.backend != "python" else ()),
        *(("Rust run used the wrong backend",) if rust.backend != "rust" else ()),
        *(("Python/Rust test inventories differ",) if python.tests != rust.tests else ()),
        *(("Python/Rust test outcomes differ",) if python_only or rust_only else ()),
        *(f"Python only: {nodeid} [{stage}] {outcome}" for nodeid, stage, outcome in python_only),
        *(f"Rust only: {nodeid} [{stage}] {outcome}" for nodeid, stage, outcome in rust_only),
        *(f"Python run: {problem}" for problem in python.problems if not python.verified or not python.tests),
        *(f"Rust run: {problem}" for problem in rust.problems if not rust.verified or not rust.tests),
        *(("no Python tests collected",) if not python.tests else ()),
        *(("Python/Rust exit codes differ",) if python.exit_code != rust.exit_code else ()),
    )


def _load_probe(reference: str) -> Callable[[], object] | None:
    if not reference:
        return None
    module, name = reference.rsplit(":", 1)
    return cast(Callable[[], object], getattr(importlib.import_module(module), name))


def main(argv: Sequence[str] | None = None) -> int:
    import pytest

    parser: Final = argparse.ArgumentParser()
    parser.add_argument("--backend", required=True, choices=("python", "rust"))
    parser.add_argument("--probe", default="")
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument("pytest_args", nargs=argparse.REMAINDER)
    namespace: Final = parser.parse_args(argv)
    args: Final = WorkerArgs.model_validate(vars(namespace))
    try:
        plugin: Final = ResultPlugin(args.backend, _load_probe(args.probe))
        plugin.verify()
        code: Final = int(
            pytest.main(["-o", "consider_namespace_packages=true", *args.pytest_args[1:]], plugins=[plugin])
        )
        report: Final = PythonReport(
            backend=args.backend,
            verified=True,
            tests=plugin.tests,
            outcomes=plugin.outcomes,
            exit_code=code,
            problems=plugin.problems,
        )
    except Exception as error:
        failure: Final = PythonReport(backend=args.backend, verified=False, exit_code=1, problems=(str(error),))
        args.output.write_text(failure.model_dump_json())
        return 1
    args.output.write_text(report.model_dump_json())
    return report.exit_code


def contract_nodeid(nodeid: str) -> str:
    owner, separator, test = nodeid.rpartition("::")
    function: Final = test.partition("[")[0]
    if not separator or not function.startswith("test_"):
        raise ValueError(f"Unrecognized pytest node id: {nodeid}")
    return f"{owner}::{function}"


def collect_python_tests(selectors: Sequence[str], repo_root: Path) -> frozenset[str]:
    report: Final = run_python_tests(
        selectors,
        repo_root,
        "python",
        BackendSpec(environment_variable="LITELLM_RUST"),
        ("--collect-only", "-p", "no:cacheprovider"),
    )
    if report.exit_code or report.problems:
        details: Final = "\n".join(report.problems) or f"pytest exited with code {report.exit_code}"
        raise ValueError(f"Python test collection failed:\n{details}")
    tests: Final = frozenset(contract_nodeid(nodeid) for nodeid in report.tests)
    if not tests:
        raise ValueError(f"pytest collected no tests for: {', '.join(selectors)}")
    return tests


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