mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-16 04:29:31 +00:00
182 lines
7.3 KiB
Python
182 lines
7.3 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import importlib.util
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from collections.abc import Mapping
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from types import MappingProxyType, ModuleType
|
|
from typing import Final
|
|
|
|
GROUPS: Final = MappingProxyType(
|
|
{
|
|
"management": ("management", "authorization", "configuration"),
|
|
"accounting": ("pricing", "spend", "caching"),
|
|
"database": ("database",),
|
|
"providers": ("providers", "routing", "streaming", "messages_endpoint", "translation"),
|
|
"extensions": ("observability", "compatibility"),
|
|
"mcp": ("mcp",),
|
|
"sdk": ("sdk",),
|
|
"cost": ("cost_calculation",),
|
|
"security": ("security",),
|
|
}
|
|
)
|
|
GITHUB_FILES: Final = frozenset(
|
|
{
|
|
"tests/integration/database/test_roi_observed.py",
|
|
"tests/integration/database/image/test_component_image_serves_offline.py",
|
|
"tests/integration/database/image/test_image_admin_mcp.py",
|
|
"tests/integration/database/image/test_image_bedrock_realtime_extra.py",
|
|
"tests/integration/database/image/test_offline_image_migration.py",
|
|
"tests/integration/database/image/test_ui_image_serves_offline.py",
|
|
}
|
|
)
|
|
|
|
|
|
PEER_GROUPS: Final = frozenset({"management", "mcp", "accounting"})
|
|
COST_MAP_GROUPS: Final = frozenset({"cost"})
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class Selection:
|
|
nodes: tuple[str, ...]
|
|
foreign: tuple[str, ...]
|
|
|
|
|
|
def file_of(node: str) -> str:
|
|
return node.split("::", 1)[0]
|
|
|
|
|
|
def select(requested: tuple[str, ...], group_files: tuple[str, ...]) -> Selection:
|
|
members: Final = frozenset(group_files)
|
|
return Selection(
|
|
nodes=requested or group_files,
|
|
foreign=tuple(sorted({node for node in requested if file_of(node) not in members})),
|
|
)
|
|
|
|
|
|
def uncollected(nodes: tuple[str, ...], collected: frozenset[str]) -> tuple[str, ...]:
|
|
collected_files: Final = frozenset(file_of(node) for node in collected)
|
|
return tuple(node for node in nodes if file_of(node) not in collected_files)
|
|
|
|
|
|
def _without_provider_api_keys(environment: Mapping[str, str]) -> dict[str, str]:
|
|
return {name: value for name, value in environment.items() if not name.endswith("_API_KEY")}
|
|
|
|
|
|
def worker_settings(group: str, workers: int) -> Mapping[str, str]:
|
|
if workers < 2:
|
|
return MappingProxyType({})
|
|
return MappingProxyType(
|
|
{
|
|
"INTEGRATION_WORKER_ISOLATION": "1",
|
|
"INTEGRATION_WORKER_PEER": "1" if group in PEER_GROUPS else "",
|
|
"INTEGRATION_WORKER_COST_MAP": "1" if group in COST_MAP_GROUPS else "",
|
|
}
|
|
)
|
|
|
|
|
|
def _budget() -> ModuleType:
|
|
specification: Final = importlib.util.spec_from_file_location(
|
|
"integration_budget", Path(__file__).with_name("_support") / "budget.py"
|
|
)
|
|
assert specification is not None and specification.loader is not None
|
|
module: Final = importlib.util.module_from_spec(specification)
|
|
specification.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
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")))
|
|
parser.add_argument("--shards", type=int, default=int(os.environ.get("INTEGRATION_CASE_SHARDS", "1")))
|
|
parser.add_argument("--shard-index", type=int, default=int(os.environ.get("INTEGRATION_CASE_SHARD_INDEX", "0")))
|
|
parser.add_argument("--shard-files", type=Path)
|
|
parser.add_argument("--list", action="store_true", help="print selected candidate files and exit")
|
|
parser.add_argument("files", nargs="*", help="run only these files, or pytest node ids inside them, of the group")
|
|
options: Final = parser.parse_intermixed_args()
|
|
root: Final = Path(__file__).resolve().parents[2]
|
|
group_files: Final = tuple(
|
|
str(path.relative_to(root))
|
|
for folder in GROUPS[options.group]
|
|
for path in sorted((root / "tests/integration" / folder).rglob("test_*.py"))
|
|
if str(path.relative_to(root)) not in GITHUB_FILES
|
|
)
|
|
selection: Final = select(tuple(options.files), group_files)
|
|
if selection.foreign:
|
|
parser.error(f"Not in the {options.group} group: {', '.join(selection.foreign)}")
|
|
if not selection.nodes:
|
|
parser.error(f"No integration test files selected for {options.group}")
|
|
if options.list:
|
|
sys.stdout.write("\n".join(sorted({file_of(node) for node in selection.nodes})) + "\n")
|
|
return 0
|
|
output: Final = options.results.resolve()
|
|
output.mkdir(parents=True, exist_ok=True)
|
|
environment: Final = {
|
|
**_without_provider_api_keys(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",
|
|
"LITELLM_MODE": "PRODUCTION",
|
|
"PYTHON_DOTENV_DISABLED": "1",
|
|
**worker_settings(options.group, options.workers),
|
|
}
|
|
budget: Final = _budget()
|
|
proxy_root: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or root)
|
|
imported: Final = float(
|
|
os.environ.get(budget.IMPORT_SECONDS_VARIABLE) or budget.measure_import_seconds(proxy_root, environment)
|
|
)
|
|
(output / "boot-calibration.json").write_text(budget.calibration(imported) + "\n")
|
|
result: Final = subprocess.call(
|
|
[
|
|
sys.executable,
|
|
"-m",
|
|
"pytest",
|
|
*selection.nodes,
|
|
"-vv",
|
|
"-rs",
|
|
"--strict-markers",
|
|
"-p",
|
|
"no:pytest-retry",
|
|
"-p",
|
|
"no:rerunfailures",
|
|
f"--timeout={int(budget.case_timeout_for(imported))}",
|
|
"--durations=15",
|
|
"--tb=short",
|
|
f"--hypothesis-seed={options.seed}",
|
|
f"--integration-order-seed={options.order_seed}",
|
|
f"--integration-shard-count={options.shards}",
|
|
f"--integration-shard-index={options.shard_index}",
|
|
*(("--integration-shard-files", str(options.shard_files)) if options.shard_files is not None else ()),
|
|
f"--junitxml={output / 'junit.xml'}",
|
|
"-o",
|
|
"junit_family=xunit1",
|
|
*(("-n", str(options.workers)) if options.workers > 1 else ()),
|
|
*(("--dist=loadgroup",) if options.workers > 1 else ()),
|
|
],
|
|
cwd=root,
|
|
env={**environment, budget.IMPORT_SECONDS_VARIABLE: str(imported)},
|
|
)
|
|
if result != 0:
|
|
return result
|
|
evidence: Final = json.loads((output / "execution.json").read_text())
|
|
empty: Final = uncollected(selection.nodes, frozenset(evidence["inventory"]))
|
|
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())
|