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 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))) 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_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