mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
208 lines
7.4 KiB
Python
208 lines
7.4 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import ast
|
|
import importlib
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
from collections.abc import Callable, Sequence
|
|
from pathlib import Path
|
|
from typing import Final, Literal, cast
|
|
|
|
import pytest
|
|
from pydantic import BaseModel, ConfigDict
|
|
|
|
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
|
|
|
|
|
|
def ocr_backend() -> Backend:
|
|
from litellm.rust_bridge import native_bridge_available
|
|
from litellm.rust_bridge.configuration import rust_ocr_enabled
|
|
|
|
if not rust_ocr_enabled():
|
|
return "python"
|
|
if not native_bridge_available():
|
|
raise RuntimeError("Rust OCR was enabled but the native extension is unavailable")
|
|
return "rust"
|
|
|
|
|
|
class ResultPlugin:
|
|
def __init__(self, backend: Backend, probe: Callable[[], object]) -> 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() != 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)
|
|
|
|
@pytest.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,
|
|
"--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())
|
|
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, ...]:
|
|
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 sorted(python.outcomes) != sorted(rust.outcomes) else ()),
|
|
*(("no Python tests collected",) if not python.tests else ()),
|
|
*(
|
|
("Python tests did not all pass",)
|
|
if set(python.tests)
|
|
!= {node for node, phase, status in python.outcomes if phase == "call" and status == "passed"}
|
|
else ()
|
|
),
|
|
*(
|
|
("Rust-enabled Python tests did not all pass",)
|
|
if set(rust.tests)
|
|
!= {node for node, phase, status in rust.outcomes if phase == "call" and status == "passed"}
|
|
else ()
|
|
),
|
|
*(("Python test run failed",) if python.exit_code else ()),
|
|
*(("Rust-enabled Python test run failed",) if rust.exit_code else ()),
|
|
*python.problems,
|
|
*rust.problems,
|
|
)
|
|
|
|
|
|
def main(argv: Sequence[str] | None = None) -> int:
|
|
parser: Final = argparse.ArgumentParser()
|
|
parser.add_argument("--backend", required=True, choices=("python", "rust"))
|
|
parser.add_argument("--probe", required=True)
|
|
parser.add_argument("--output", required=True, type=Path)
|
|
parser.add_argument("pytest_args", nargs=argparse.REMAINDER)
|
|
args: Final = parser.parse_args(argv)
|
|
try:
|
|
module, name = args.probe.rsplit(":", 1)
|
|
probe: Final = cast(Callable[[], object], getattr(importlib.import_module(module), name))
|
|
plugin: Final = ResultPlugin(args.backend, 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 enumerate_python_tests(repo_root: Path, relative_path: str) -> frozenset[str]:
|
|
source = (repo_root / relative_path).read_text(encoding="utf-8")
|
|
tree = ast.parse(source, filename=relative_path)
|
|
|
|
module_level: list[str] = []
|
|
for node in ast.iter_child_nodes(tree):
|
|
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name.startswith("test_"):
|
|
module_level.append(node.name)
|
|
elif isinstance(node, ast.ClassDef):
|
|
for child in ast.iter_child_nodes(node):
|
|
if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)) and child.name.startswith(
|
|
"test_"
|
|
):
|
|
module_level.append(f"{node.name}::{child.name}")
|
|
|
|
return frozenset(module_level)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|