mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
273 lines
12 KiB
Python
273 lines
12 KiB
Python
import importlib.util
|
|
import json
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
from types import ModuleType
|
|
from typing import Final
|
|
|
|
import pytest
|
|
|
|
ROOT: Final = Path(__file__).resolve().parents[2]
|
|
CI_CD: Final = ROOT / "ci_cd"
|
|
|
|
|
|
def _load(name: str) -> ModuleType:
|
|
spec = importlib.util.spec_from_file_location(name, CI_CD / f"{name}.py")
|
|
assert spec is not None and spec.loader is not None
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
schema_module: Final = _load("generate_model_prices_schema")
|
|
guard: Final = _load("cost_map_guard")
|
|
|
|
MAP_FILES: Final = (guard.COST_MAP_PATH,)
|
|
BOT_REF: Final = "litellm_cost_map_sync_2026-09-04T12-00Z"
|
|
|
|
|
|
def _entry(price: float = 1e-06, **extra: object) -> dict[str, object]:
|
|
return {
|
|
"input_cost_per_token": price,
|
|
"output_cost_per_token": price * 2,
|
|
"litellm_provider": "openrouter",
|
|
"mode": "chat",
|
|
"max_tokens": 4096,
|
|
**extra,
|
|
}
|
|
|
|
|
|
BASE_MAP: Final = {
|
|
"sample_spec": {"input_cost_per_token": "USD per prompt token"},
|
|
"fallback_generalizations": {"rules": [{"name": "r", "pattern": "^x"}]},
|
|
"openrouter/a": _entry(supports_vision=True),
|
|
"openrouter/b": _entry(2e-06),
|
|
}
|
|
|
|
|
|
def _serialize(cost_map: dict[str, object]) -> str:
|
|
return json.dumps(cost_map, indent=4, ensure_ascii=False) + "\n"
|
|
|
|
|
|
def _snapshot(cost_map: dict[str, object], backup: str | None = None, schema: str | None = None) -> object:
|
|
text = _serialize(cost_map)
|
|
rendered = schema_module.render(schema_module.build_schema(cost_map))
|
|
return guard.Snapshot(
|
|
cost_map=text, backup=text if backup is None else backup, schema=rendered if schema is None else schema
|
|
)
|
|
|
|
|
|
BASE: Final = _snapshot(BASE_MAP)
|
|
|
|
|
|
def _failures(head: object, changed_files: tuple[str, ...] = MAP_FILES, bot: bool = True) -> tuple[str, ...]:
|
|
return guard.guard_failures(BASE, head, changed_files, bot)
|
|
|
|
|
|
def test_in_sync_files_pass_for_humans_and_bots() -> None:
|
|
assert _failures(BASE, bot=False) == ()
|
|
assert _failures(BASE, bot=True) == ()
|
|
|
|
|
|
def test_bot_may_add_and_reprice_models() -> None:
|
|
head = _snapshot({**BASE_MAP, "openrouter/a": _entry(9e-06, supports_vision=True), "openrouter/c": _entry()})
|
|
assert _failures(head) == ()
|
|
|
|
|
|
def test_broken_json_is_reported() -> None:
|
|
head = guard.Snapshot(cost_map="{not json", backup="{not json", schema="{}")
|
|
assert _failures(head, bot=False) == (
|
|
f"{guard.COST_MAP_PATH} is not valid JSON: Expecting property name enclosed in double quotes: line 1 column 2 (char 1)",
|
|
)
|
|
|
|
|
|
def test_non_object_root_is_reported() -> None:
|
|
head = guard.Snapshot(cost_map="[]", backup="[]", schema="{}")
|
|
assert _failures(head, bot=False) == (f"{guard.COST_MAP_PATH} must be a JSON object at the root",)
|
|
|
|
|
|
def test_backup_drift_is_reported() -> None:
|
|
head = _snapshot(BASE_MAP, backup=_serialize({**BASE_MAP, "openrouter/b": _entry(3e-06)}))
|
|
assert [failure for failure in _failures(head, bot=False) if failure.startswith(guard.BACKUP_PATH)]
|
|
|
|
|
|
def test_schema_out_of_sync_is_reported() -> None:
|
|
head = _snapshot({**BASE_MAP, "openrouter/c": _entry(supports_audio_input=True)}, schema=BASE.schema)
|
|
assert [failure for failure in _failures(head, bot=False) if failure.startswith(guard.SCHEMA_PATH)]
|
|
|
|
|
|
def test_schema_validation_errors_are_reported() -> None:
|
|
head = _snapshot({**BASE_MAP, "openrouter/c": _entry(-1e-06)})
|
|
prefix = f"{guard.COST_MAP_PATH} does not validate against its schema: openrouter/c."
|
|
assert [failure.removeprefix(prefix).split(":")[0] for failure in _failures(head, bot=False)] == [
|
|
"input_cost_per_token",
|
|
"output_cost_per_token",
|
|
]
|
|
|
|
|
|
def test_unclassified_entry_key_is_reported() -> None:
|
|
text = _serialize({**BASE_MAP, "openrouter/c": _entry(weird_thing=1)})
|
|
head = guard.Snapshot(cost_map=text, backup=text, schema=BASE.schema)
|
|
(failure,) = _failures(head, bot=False)
|
|
assert "Unclassified keys" in failure and "weird_thing" in failure
|
|
|
|
|
|
STALE_HEAD: Final = _snapshot(BASE_MAP, backup=_serialize({**BASE_MAP, "openrouter/b": _entry(3e-06)}), schema="{}")
|
|
CODE_ONLY: Final = (
|
|
"litellm/utils.py",
|
|
"tests/test_litellm/test_utils.py",
|
|
"docs/model_prices_and_context_window.json",
|
|
)
|
|
|
|
|
|
def test_human_pr_that_leaves_the_cost_map_alone_skips_the_file_checks() -> None:
|
|
unparseable: Final = guard.Snapshot(cost_map="{not json", backup="", schema="")
|
|
assert _failures(STALE_HEAD, changed_files=CODE_ONLY, bot=False) == ()
|
|
assert _failures(STALE_HEAD, changed_files=(), bot=False) == ()
|
|
assert _failures(unparseable, changed_files=CODE_ONLY, bot=False) == ()
|
|
|
|
|
|
@pytest.mark.parametrize("guarded_path", guard.GUARDED_PATHS)
|
|
def test_touching_any_cost_map_file_keeps_the_file_checks(guarded_path: str) -> None:
|
|
failures: Final = _failures(STALE_HEAD, changed_files=(*CODE_ONLY, guarded_path), bot=False)
|
|
assert [failure for failure in failures if failure.startswith(guard.BACKUP_PATH)]
|
|
assert [failure for failure in failures if failure.startswith(guard.SCHEMA_PATH)]
|
|
|
|
|
|
def test_bot_pr_always_gets_the_file_checks() -> None:
|
|
failures: Final = _failures(STALE_HEAD, changed_files=CODE_ONLY, bot=True)
|
|
assert [failure for failure in failures if failure.startswith(guard.BACKUP_PATH)]
|
|
assert [failure for failure in failures if failure.startswith(guard.SCHEMA_PATH)]
|
|
|
|
|
|
def test_contract_names_the_skip() -> None:
|
|
assert guard.contract_for(False, CODE_ONLY) == "human PR, cost map untouched"
|
|
assert guard.contract_for(False, (*CODE_ONLY, guard.SCHEMA_PATH)) == "human PR, file checks only"
|
|
assert guard.contract_for(True, CODE_ONLY) == "bot contract enforced"
|
|
|
|
|
|
def test_bot_may_only_touch_the_cost_map_files() -> None:
|
|
changed = (*guard.GUARDED_PATHS, "litellm/utils.py", ".github/workflows/cost-map-guard.yml")
|
|
assert _failures(BASE, changed_files=changed, bot=False) == ()
|
|
assert _failures(BASE, changed_files=changed) == (
|
|
"bot PRs may only change the cost map files, not litellm/utils.py",
|
|
"bot PRs may only change the cost map files, not .github/workflows/cost-map-guard.yml",
|
|
)
|
|
|
|
|
|
def test_bot_may_not_remove_models() -> None:
|
|
head = _snapshot({key: value for key, value in BASE_MAP.items() if key != "openrouter/b"})
|
|
assert _failures(head, bot=False) == ()
|
|
assert _failures(head) == ("bot PRs may not remove models: openrouter/b",)
|
|
|
|
|
|
def test_bot_may_not_remove_fields() -> None:
|
|
head = _snapshot({**BASE_MAP, "openrouter/a": _entry()})
|
|
assert _failures(head, bot=False) == ()
|
|
assert _failures(head) == ("bot PRs may not remove fields: openrouter/a.supports_vision",)
|
|
|
|
|
|
def test_bot_may_not_change_special_root_keys() -> None:
|
|
head = _snapshot({**BASE_MAP, "fallback_generalizations": {"rules": []}})
|
|
assert _failures(head, bot=False) == ()
|
|
assert _failures(head) == ("bot PRs may not change fallback_generalizations",)
|
|
|
|
|
|
def _commit(repo: Path, cost_map: dict[str, object], message: str) -> str:
|
|
text = _serialize(cost_map)
|
|
(repo / guard.COST_MAP_PATH).write_text(text)
|
|
(repo / guard.BACKUP_PATH).parent.mkdir(exist_ok=True)
|
|
(repo / guard.BACKUP_PATH).write_text(text)
|
|
(repo / guard.SCHEMA_PATH).write_text(schema_module.render(schema_module.build_schema(cost_map)))
|
|
return _git_commit(repo, message)
|
|
|
|
|
|
def _commit_code_only(repo: Path, message: str) -> str:
|
|
(repo / "litellm").mkdir(exist_ok=True)
|
|
(repo / "litellm" / "utils.py").write_text(f"print('{message}')\n")
|
|
return _git_commit(repo, message)
|
|
|
|
|
|
def _git_commit(repo: Path, message: str) -> str:
|
|
subprocess.run(("git", "add", "-A"), cwd=repo, check=True)
|
|
subprocess.run(
|
|
("git", "-c", "user.name=t", "-c", "user.email=t@example.com", "commit", "-q", "-m", message),
|
|
cwd=repo,
|
|
check=True,
|
|
)
|
|
return subprocess.run(
|
|
("git", "rev-parse", "HEAD"), cwd=repo, check=True, capture_output=True, text=True
|
|
).stdout.strip()
|
|
|
|
|
|
def _run_guard(repo: Path, base: str, head: str, head_ref: str) -> subprocess.CompletedProcess[str]:
|
|
return subprocess.run(
|
|
(sys.executable, str(CI_CD / "cost_map_guard.py"), "--base", base, "--head", head, "--head-ref", head_ref),
|
|
cwd=repo,
|
|
capture_output=True,
|
|
text=True,
|
|
check=False,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("head_ref", "expected_code", "expected_line"),
|
|
[
|
|
(BOT_REF, 1, "- bot PRs may not remove models: openrouter/b"),
|
|
("litellm_fix_pricing", 0, "cost map guard passed (human PR, file checks only)"),
|
|
],
|
|
)
|
|
def test_main_reads_both_revisions_from_git(
|
|
tmp_path: Path, head_ref: str, expected_code: int, expected_line: str
|
|
) -> None:
|
|
subprocess.run(("git", "init", "-q", str(tmp_path)), check=True)
|
|
base = _commit(tmp_path, BASE_MAP, "base")
|
|
head = _commit(tmp_path, {key: value for key, value in BASE_MAP.items() if key != "openrouter/b"}, "head")
|
|
result = _run_guard(tmp_path, base, head, head_ref)
|
|
assert result.returncode == expected_code, result.stdout + result.stderr
|
|
assert expected_line in result.stdout.splitlines()
|
|
|
|
|
|
def test_main_skips_the_file_checks_on_a_stale_base_the_pr_never_touched(tmp_path: Path) -> None:
|
|
subprocess.run(("git", "init", "-q", str(tmp_path)), check=True)
|
|
_commit(tmp_path, BASE_MAP, "base")
|
|
(tmp_path / guard.BACKUP_PATH).write_text(_serialize({**BASE_MAP, "openrouter/b": _entry(3e-06)}))
|
|
(tmp_path / guard.SCHEMA_PATH).write_text("{}")
|
|
stale_base: Final = _commit_code_only(tmp_path, "stale base with drifted backup and schema")
|
|
head: Final = _commit_code_only(tmp_path, "code change on the stale base")
|
|
human: Final = _run_guard(tmp_path, stale_base, head, "litellm_fix_pricing")
|
|
assert human.returncode == 0, human.stdout + human.stderr
|
|
assert "cost map guard passed (human PR, cost map untouched)" in human.stdout.splitlines()
|
|
bot: Final = _run_guard(tmp_path, stale_base, head, BOT_REF)
|
|
assert bot.returncode == 1
|
|
backup_failure: Final = f"- {guard.BACKUP_PATH} differs from {guard.COST_MAP_PATH}; copy the root file over it"
|
|
assert backup_failure in bot.stdout.splitlines()
|
|
|
|
|
|
def test_main_keeps_the_file_checks_when_a_cost_map_file_is_renamed(tmp_path: Path) -> None:
|
|
subprocess.run(("git", "init", "-q", str(tmp_path)), check=True)
|
|
base: Final = _commit(tmp_path, BASE_MAP, "base")
|
|
subprocess.run(("git", "mv", guard.COST_MAP_PATH, "renamed.json"), cwd=tmp_path, check=True)
|
|
head: Final = _git_commit(tmp_path, "rename the cost map")
|
|
result: Final = _run_guard(tmp_path, base, head, "litellm_fix_pricing")
|
|
assert result.returncode == 1, result.stdout + result.stderr
|
|
assert "cost map guard failed (human PR, file checks only):" in result.stdout.splitlines()
|
|
|
|
|
|
def test_main_fails_when_the_changed_files_cannot_be_read(tmp_path: Path) -> None:
|
|
subprocess.run(("git", "init", "-q", str(tmp_path)), check=True)
|
|
head: Final = _commit(tmp_path, BASE_MAP, "head")
|
|
result: Final = _run_guard(tmp_path, "0" * 40, head, "litellm_fix_pricing")
|
|
assert result.returncode == 1, result.stdout + result.stderr
|
|
assert result.stdout.startswith("cost map guard failed: git diff ")
|
|
|
|
|
|
def test_main_rejects_a_bot_pr_that_edits_code(tmp_path: Path) -> None:
|
|
subprocess.run(("git", "init", "-q", str(tmp_path)), check=True)
|
|
base = _commit(tmp_path, BASE_MAP, "base")
|
|
(tmp_path / "litellm" / "utils.py").write_text("print('hi')\n")
|
|
head = _commit(tmp_path, {**BASE_MAP, "openrouter/c": _entry()}, "head")
|
|
assert _run_guard(tmp_path, base, head, BOT_REF).returncode == 1
|
|
assert _run_guard(tmp_path, base, head, "litellm_fix_pricing").returncode == 0
|