litellm/tests/unit/test_cost_map_guard.py

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