From 4576f9cabadd062bd5d7ec78e00de37d8b087fb6 Mon Sep 17 00:00:00 2001 From: JerrettDavis Date: Wed, 22 Apr 2026 22:13:13 -0500 Subject: [PATCH] test: remove provider diff churn Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- tests/test_install/test_providers.py | 936 +++++++++++------------ tests/test_provider_cursor.py | 60 +- tests/test_provider_registry.py | 786 ++++++++++---------- tests/test_providers/test_universal.py | 990 ++++++++++++------------- 4 files changed, 1386 insertions(+), 1386 deletions(-) diff --git a/tests/test_install/test_providers.py b/tests/test_install/test_providers.py index 58f79453c..5c9fba42c 100644 --- a/tests/test_install/test_providers.py +++ b/tests/test_install/test_providers.py @@ -1,468 +1,468 @@ -from __future__ import annotations - -import json -import os -from pathlib import Path - -import click -import pytest - -from headroom.install.models import DeploymentManifest, ManagedMutation -from headroom.install.providers import _apply_windows_env_scope, _remove_windows_env_scope -from headroom.providers.claude.install import apply_provider_scope as apply_claude_provider_scope -from headroom.providers.claude.install import build_install_env as build_claude_install_env -from headroom.providers.claude.install import revert_provider_scope as revert_claude_provider_scope -from headroom.providers.codex.install import apply_provider_scope as apply_codex_provider_scope -from headroom.providers.codex.install import build_install_env as build_codex_install_env -from headroom.providers.codex.install import revert_provider_scope as revert_codex_provider_scope -from headroom.providers.copilot.install import build_install_env as build_copilot_install_env - - -def _manifest(tmp_path: Path) -> DeploymentManifest: - return DeploymentManifest( - profile="default", - preset="persistent-service", - runtime_kind="python", - supervisor_kind="service", - scope="provider", - provider_mode="manual", - targets=["claude", "codex"], - port=8787, - host="127.0.0.1", - backend="anthropic", - memory_db_path=str(tmp_path / "memory.db"), - tool_envs={ - "claude": {"ANTHROPIC_BASE_URL": "http://127.0.0.1:8787"}, - "codex": {"OPENAI_BASE_URL": "http://127.0.0.1:8787/v1"}, - }, - ) - - -def test_apply_and_revert_claude_provider_scope(monkeypatch, tmp_path: Path) -> None: - settings_path = tmp_path / "settings.json" - settings_path.write_text( - json.dumps({"env": {"ANTHROPIC_API_KEY": "keep", "ANTHROPIC_BASE_URL": "https://old"}}) - ) - monkeypatch.setattr( - "headroom.providers.claude.install.claude_settings_path", lambda: settings_path - ) - manifest = _manifest(tmp_path) - - mutation = apply_claude_provider_scope(manifest) - payload = json.loads(settings_path.read_text()) - assert payload["env"]["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:8787" - assert payload["env"]["ANTHROPIC_API_KEY"] == "keep" - - assert mutation is not None - revert_claude_provider_scope(mutation, manifest) - reverted = json.loads(settings_path.read_text()) - assert reverted["env"]["ANTHROPIC_BASE_URL"] == "https://old" - assert reverted["env"]["ANTHROPIC_API_KEY"] == "keep" - - -def test_apply_and_revert_codex_provider_scope(monkeypatch, tmp_path: Path) -> None: - config_path = tmp_path / "config.toml" - config_path.write_text('model = "gpt-4o"\n') - monkeypatch.setattr("headroom.providers.codex.install.codex_config_path", lambda: config_path) - manifest = _manifest(tmp_path) - - mutation = apply_codex_provider_scope(manifest) - content = config_path.read_text() - assert 'model_provider = "headroom"' in content - assert 'base_url = "http://127.0.0.1:8787/v1"' in content - - assert mutation is not None - revert_codex_provider_scope(mutation, manifest) - reverted = config_path.read_text() - assert 'model_provider = "headroom"' not in reverted - assert reverted.strip() == 'model = "gpt-4o"' - - -def test_codex_build_install_env_returns_proxy_base_url() -> None: - env = build_codex_install_env(port=5566, backend="ignored") - - assert env == {"OPENAI_BASE_URL": "http://127.0.0.1:5566/v1"} - - -def test_apply_codex_provider_scope_skips_non_provider_scope(monkeypatch, tmp_path: Path) -> None: - config_path = tmp_path / "config.toml" - monkeypatch.setattr("headroom.providers.codex.install.codex_config_path", lambda: config_path) - manifest = _manifest(tmp_path) - manifest.scope = "user" - - mutation = apply_codex_provider_scope(manifest) - - assert mutation is None - assert not config_path.exists() - - -def test_apply_codex_provider_scope_replaces_existing_managed_block( - monkeypatch, tmp_path: Path -) -> None: - config_path = tmp_path / "config.toml" - config_path.write_text( - 'model = "gpt-4o"\n\n' - "# --- Headroom persistent provider ---\n" - 'model_provider = "headroom"\n\n' - "[model_providers.headroom]\n" - 'name = "Headroom persistent proxy"\n' - 'base_url = "http://127.0.0.1:1111/v1"\n' - 'env_key = "OPENAI_API_KEY"\n' - "requires_openai_auth = true\n" - "supports_websockets = true\n" - "# --- end Headroom persistent provider ---\n" - ) - monkeypatch.setattr("headroom.providers.codex.install.codex_config_path", lambda: config_path) - manifest = _manifest(tmp_path) - manifest.port = 9999 - - apply_codex_provider_scope(manifest) - - content = config_path.read_text() - assert content.count("# --- Headroom persistent provider ---") == 1 - assert 'base_url = "http://127.0.0.1:9999/v1"' in content - assert 'base_url = "http://127.0.0.1:1111/v1"' not in content - - -def test_apply_codex_provider_scope_creates_new_config_when_missing( - monkeypatch, tmp_path: Path -) -> None: - config_path = tmp_path / "nested" / "config.toml" - monkeypatch.setattr("headroom.providers.codex.install.codex_config_path", lambda: config_path) - manifest = _manifest(tmp_path) - - mutation = apply_codex_provider_scope(manifest) - - assert mutation is not None - assert 'base_url = "http://127.0.0.1:8787/v1"' in config_path.read_text() - - -def test_revert_codex_provider_scope_ignores_missing_path_and_file(tmp_path: Path) -> None: - manifest = _manifest(tmp_path) - - revert_codex_provider_scope( - ManagedMutation(target="codex", kind="toml-block"), - manifest, - ) - revert_codex_provider_scope( - ManagedMutation( - target="codex", - kind="toml-block", - path=str(tmp_path / "missing.toml"), - ), - manifest, - ) - - -def test_revert_codex_provider_scope_ignores_files_without_managed_block( - monkeypatch, tmp_path: Path -) -> None: - config_path = tmp_path / "config.toml" - config_path.write_text('model = "gpt-4o"\n') - monkeypatch.setattr("headroom.providers.codex.install.codex_config_path", lambda: config_path) - manifest = _manifest(tmp_path) - mutation = ManagedMutation(target="codex", kind="toml-block", path=str(config_path)) - - revert_codex_provider_scope(mutation, manifest) - - assert config_path.read_text() == 'model = "gpt-4o"\n' - - -def test_apply_openclaw_provider_scope_uses_manifest_port(monkeypatch, tmp_path: Path) -> None: - recorded: list[list[str]] = [] - monkeypatch.setattr("headroom.providers.openclaw.install.shutil_which", lambda name: "openclaw") - monkeypatch.setattr( - "headroom.providers.openclaw.install.resolve_headroom_command", - lambda: ["headroom"], - ) - monkeypatch.setattr( - "headroom.providers.openclaw.install._invoke_openclaw", - lambda command: recorded.append(command), - ) - monkeypatch.setattr( - "headroom.providers.openclaw.install.openclaw_config_path", - lambda: tmp_path / "openclaw.json", - ) - manifest = _manifest(tmp_path) - manifest.port = 9999 - - from headroom.providers.openclaw.install import ( - apply_provider_scope as apply_openclaw_provider_scope, - ) - - apply_openclaw_provider_scope(manifest) - - assert recorded == [["headroom", "wrap", "openclaw", "--no-auto-start", "--proxy-port", "9999"]] - - -def test_openclaw_apply_provider_scope_requires_installed_binary( - tmp_path: Path, monkeypatch -) -> None: - monkeypatch.setattr("headroom.providers.openclaw.install.shutil_which", lambda name: None) - - with pytest.raises(click.ClickException, match="openclaw not found"): - from headroom.providers.openclaw.install import ( - apply_provider_scope as apply_openclaw_provider_scope, - ) - - apply_openclaw_provider_scope(_manifest(tmp_path)) - - -def test_openclaw_helper_wrappers_delegate_to_stdlib(monkeypatch) -> None: - monkeypatch.setattr("shutil.which", lambda name: f"/fake/{name}") - recorded: list[tuple[list[str], bool]] = [] - - def fake_run(command: list[str], check: bool) -> None: - recorded.append((command, check)) - - monkeypatch.setattr("subprocess.run", fake_run) - - from headroom.providers.openclaw.install import _invoke_openclaw, shutil_which - - assert shutil_which("openclaw") == "/fake/openclaw" - _invoke_openclaw(["headroom", "wrap", "openclaw"]) - - assert recorded == [(["headroom", "wrap", "openclaw"], True)] - - -def test_openclaw_revert_provider_scope_skips_without_binary(monkeypatch, tmp_path: Path) -> None: - monkeypatch.setattr("headroom.providers.openclaw.install.shutil_which", lambda name: None) - called = False - - def fail_if_called(command: list[str]) -> None: - nonlocal called - called = True - - monkeypatch.setattr("headroom.providers.openclaw.install._invoke_openclaw", fail_if_called) - - from headroom.providers.openclaw.install import ( - revert_provider_scope as revert_openclaw_provider_scope, - ) - - revert_openclaw_provider_scope( - ManagedMutation(target="openclaw", kind="openclaw-wrap", path=str(tmp_path / "cfg.json")), - _manifest(tmp_path), - ) - - assert called is False - - -def test_openclaw_revert_provider_scope_invokes_unwrap(monkeypatch, tmp_path: Path) -> None: - recorded: list[list[str]] = [] - monkeypatch.setattr("headroom.providers.openclaw.install.shutil_which", lambda name: "openclaw") - monkeypatch.setattr( - "headroom.providers.openclaw.install.resolve_headroom_command", - lambda: ["headroom"], - ) - monkeypatch.setattr( - "headroom.providers.openclaw.install._invoke_openclaw", - lambda command: recorded.append(command), - ) - - from headroom.providers.openclaw.install import ( - revert_provider_scope as revert_openclaw_provider_scope, - ) - - revert_openclaw_provider_scope( - ManagedMutation(target="openclaw", kind="openclaw-wrap", path=str(tmp_path / "cfg.json")), - _manifest(tmp_path), - ) - - assert recorded == [["headroom", "unwrap", "openclaw"]] - - -def test_windows_env_scope_restores_previous_values(monkeypatch, tmp_path: Path) -> None: - manifest = _manifest(tmp_path) - manifest.scope = "user" - manifest.targets = ["claude"] - manifest.base_env = {"HEADROOM_PORT": "8787"} - manifest.tool_envs = {"claude": {"ANTHROPIC_BASE_URL": "http://127.0.0.1:8787"}} - - calls: list[list[str]] = [] - previous_values = { - "HEADROOM_PORT": "7777", - "ANTHROPIC_BASE_URL": "https://old", - } - - class Result: - def __init__(self, stdout: str = "") -> None: - self.stdout = stdout - - def fake_run(command: list[str], **kwargs): - calls.append(command) - script = command[-1] - if "GetEnvironmentVariable" in script: - name = script.split("GetEnvironmentVariable('", 1)[1].split("'", 1)[0] - value = previous_values.get(name, "__HEADROOM_UNSET__") - return Result(stdout=value) - return Result() - - monkeypatch.setattr("headroom.install.providers.subprocess.run", fake_run) - - mutations = _apply_windows_env_scope(manifest) - _remove_windows_env_scope(mutations) - - previous_by_name = {mutation.data["name"]: mutation.data["previous"] for mutation in mutations} - assert previous_by_name["HEADROOM_PORT"] == "7777" - assert previous_by_name["ANTHROPIC_BASE_URL"] == "https://old" - assert any( - "[Environment]::SetEnvironmentVariable('HEADROOM_PORT','7777','User')" in command[-1] - for command in calls - ) - assert any( - "[Environment]::SetEnvironmentVariable('ANTHROPIC_BASE_URL','https://old','User')" - in command[-1] - for command in calls - ) - - -def test_remove_windows_env_scope_requires_name_and_scope() -> None: - try: - _remove_windows_env_scope([ManagedMutation(target="env", kind="windows-env", data={})]) - except ValueError as exc: - assert "variable name" in str(exc) - else: - raise AssertionError("expected missing variable name to raise") - - try: - _remove_windows_env_scope( - [ManagedMutation(target="env", kind="windows-env", data={"name": "X", "scope": 1})] - ) - except ValueError as exc: - assert "valid scope" in str(exc) - else: - raise AssertionError("expected invalid scope to raise") - - -def test_apply_mutations_runs_openclaw_for_user_scope(monkeypatch, tmp_path: Path) -> None: - manifest = _manifest(tmp_path) - manifest.scope = "user" - manifest.targets = ["openclaw"] - manifest.base_env = {"HEADROOM_PORT": "8787"} - manifest.tool_envs = {} - - if os.name == "nt": - monkeypatch.setattr( - "headroom.install.providers._apply_windows_env_scope", lambda deployment: [] - ) - else: - monkeypatch.setattr( - "headroom.install.providers._apply_unix_env_scope", lambda deployment: [] - ) - monkeypatch.setattr( - "headroom.install.providers.apply_provider_scope_mutations", - lambda deployment: [ManagedMutation(target="openclaw", kind="openclaw-wrap")], - ) - - from headroom.install.providers import apply_mutations - - mutations = apply_mutations(manifest) - - assert [mutation.kind for mutation in mutations] == ["openclaw-wrap"] - - -def test_claude_build_install_env_returns_proxy_base_url() -> None: - # Arrange / Act - env = build_claude_install_env(port=5566, backend="ignored") - - # Assert - assert env == {"ANTHROPIC_BASE_URL": "http://127.0.0.1:5566"} - - -def test_copilot_build_install_env_uses_provider_type_specific_proxy_urls() -> None: - anthropic_env = build_copilot_install_env(port=8787, backend="anthropic") - openai_env = build_copilot_install_env(port=8787, backend="anyllm") - - assert anthropic_env == { - "COPILOT_PROVIDER_TYPE": "anthropic", - "COPILOT_PROVIDER_BASE_URL": "http://127.0.0.1:8787", - } - assert openai_env == { - "COPILOT_PROVIDER_TYPE": "openai", - "COPILOT_PROVIDER_BASE_URL": "http://127.0.0.1:8787/v1", - "COPILOT_PROVIDER_WIRE_API": "completions", - } - - -def test_apply_claude_provider_scope_skips_non_provider_scope(monkeypatch, tmp_path: Path) -> None: - # Arrange - settings_path = tmp_path / "settings.json" - monkeypatch.setattr( - "headroom.providers.claude.install.claude_settings_path", lambda: settings_path - ) - manifest = _manifest(tmp_path) - manifest.scope = "user" - - # Act - mutation = apply_claude_provider_scope(manifest) - - # Assert - assert mutation is None - assert not settings_path.exists() - - -def test_revert_claude_provider_scope_removes_new_values_from_non_mapping_env( - monkeypatch, tmp_path: Path -) -> None: - # Arrange - settings_path = tmp_path / "settings.json" - settings_path.write_text(json.dumps({"env": ["not-a-map"]})) - monkeypatch.setattr( - "headroom.providers.claude.install.claude_settings_path", lambda: settings_path - ) - manifest = _manifest(tmp_path) - - # Act - mutation = apply_claude_provider_scope(manifest) - apply_payload = json.loads(settings_path.read_text()) - revert_claude_provider_scope(mutation, manifest) - reverted_payload = json.loads(settings_path.read_text()) - - # Assert - assert mutation is not None - assert mutation.data["previous"] == {"ANTHROPIC_BASE_URL": None} - assert apply_payload["env"] == {"ANTHROPIC_BASE_URL": "http://127.0.0.1:8787"} - assert reverted_payload["env"] == {} - - -def test_apply_claude_provider_scope_creates_settings_when_missing( - monkeypatch, tmp_path: Path -) -> None: - # Arrange - settings_path = tmp_path / "nested" / "settings.json" - monkeypatch.setattr( - "headroom.providers.claude.install.claude_settings_path", lambda: settings_path - ) - manifest = _manifest(tmp_path) - - # Act - mutation = apply_claude_provider_scope(manifest) - - # Assert - assert mutation is not None - assert json.loads(settings_path.read_text()) == { - "env": {"ANTHROPIC_BASE_URL": "http://127.0.0.1:8787"} - } - - -def test_revert_claude_provider_scope_ignores_missing_mutation_path(tmp_path: Path) -> None: - # Arrange - manifest = _manifest(tmp_path) - mutation = ManagedMutation(target="claude", kind="json-env", data={"previous": {}}) - - # Act / Assert - revert_claude_provider_scope(mutation, manifest) - - -def test_revert_claude_provider_scope_ignores_missing_settings_file(tmp_path: Path) -> None: - # Arrange - manifest = _manifest(tmp_path) - mutation = ManagedMutation( - target="claude", - kind="json-env", - path=str(tmp_path / "missing-settings.json"), - data={"previous": {}}, - ) - - # Act / Assert - revert_claude_provider_scope(mutation, manifest) +from __future__ import annotations + +import json +import os +from pathlib import Path + +import click +import pytest + +from headroom.install.models import DeploymentManifest, ManagedMutation +from headroom.install.providers import _apply_windows_env_scope, _remove_windows_env_scope +from headroom.providers.claude.install import apply_provider_scope as apply_claude_provider_scope +from headroom.providers.claude.install import build_install_env as build_claude_install_env +from headroom.providers.claude.install import revert_provider_scope as revert_claude_provider_scope +from headroom.providers.codex.install import apply_provider_scope as apply_codex_provider_scope +from headroom.providers.codex.install import build_install_env as build_codex_install_env +from headroom.providers.codex.install import revert_provider_scope as revert_codex_provider_scope +from headroom.providers.copilot.install import build_install_env as build_copilot_install_env + + +def _manifest(tmp_path: Path) -> DeploymentManifest: + return DeploymentManifest( + profile="default", + preset="persistent-service", + runtime_kind="python", + supervisor_kind="service", + scope="provider", + provider_mode="manual", + targets=["claude", "codex"], + port=8787, + host="127.0.0.1", + backend="anthropic", + memory_db_path=str(tmp_path / "memory.db"), + tool_envs={ + "claude": {"ANTHROPIC_BASE_URL": "http://127.0.0.1:8787"}, + "codex": {"OPENAI_BASE_URL": "http://127.0.0.1:8787/v1"}, + }, + ) + + +def test_apply_and_revert_claude_provider_scope(monkeypatch, tmp_path: Path) -> None: + settings_path = tmp_path / "settings.json" + settings_path.write_text( + json.dumps({"env": {"ANTHROPIC_API_KEY": "keep", "ANTHROPIC_BASE_URL": "https://old"}}) + ) + monkeypatch.setattr( + "headroom.providers.claude.install.claude_settings_path", lambda: settings_path + ) + manifest = _manifest(tmp_path) + + mutation = apply_claude_provider_scope(manifest) + payload = json.loads(settings_path.read_text()) + assert payload["env"]["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:8787" + assert payload["env"]["ANTHROPIC_API_KEY"] == "keep" + + assert mutation is not None + revert_claude_provider_scope(mutation, manifest) + reverted = json.loads(settings_path.read_text()) + assert reverted["env"]["ANTHROPIC_BASE_URL"] == "https://old" + assert reverted["env"]["ANTHROPIC_API_KEY"] == "keep" + + +def test_apply_and_revert_codex_provider_scope(monkeypatch, tmp_path: Path) -> None: + config_path = tmp_path / "config.toml" + config_path.write_text('model = "gpt-4o"\n') + monkeypatch.setattr("headroom.providers.codex.install.codex_config_path", lambda: config_path) + manifest = _manifest(tmp_path) + + mutation = apply_codex_provider_scope(manifest) + content = config_path.read_text() + assert 'model_provider = "headroom"' in content + assert 'base_url = "http://127.0.0.1:8787/v1"' in content + + assert mutation is not None + revert_codex_provider_scope(mutation, manifest) + reverted = config_path.read_text() + assert 'model_provider = "headroom"' not in reverted + assert reverted.strip() == 'model = "gpt-4o"' + + +def test_codex_build_install_env_returns_proxy_base_url() -> None: + env = build_codex_install_env(port=5566, backend="ignored") + + assert env == {"OPENAI_BASE_URL": "http://127.0.0.1:5566/v1"} + + +def test_apply_codex_provider_scope_skips_non_provider_scope(monkeypatch, tmp_path: Path) -> None: + config_path = tmp_path / "config.toml" + monkeypatch.setattr("headroom.providers.codex.install.codex_config_path", lambda: config_path) + manifest = _manifest(tmp_path) + manifest.scope = "user" + + mutation = apply_codex_provider_scope(manifest) + + assert mutation is None + assert not config_path.exists() + + +def test_apply_codex_provider_scope_replaces_existing_managed_block( + monkeypatch, tmp_path: Path +) -> None: + config_path = tmp_path / "config.toml" + config_path.write_text( + 'model = "gpt-4o"\n\n' + "# --- Headroom persistent provider ---\n" + 'model_provider = "headroom"\n\n' + "[model_providers.headroom]\n" + 'name = "Headroom persistent proxy"\n' + 'base_url = "http://127.0.0.1:1111/v1"\n' + 'env_key = "OPENAI_API_KEY"\n' + "requires_openai_auth = true\n" + "supports_websockets = true\n" + "# --- end Headroom persistent provider ---\n" + ) + monkeypatch.setattr("headroom.providers.codex.install.codex_config_path", lambda: config_path) + manifest = _manifest(tmp_path) + manifest.port = 9999 + + apply_codex_provider_scope(manifest) + + content = config_path.read_text() + assert content.count("# --- Headroom persistent provider ---") == 1 + assert 'base_url = "http://127.0.0.1:9999/v1"' in content + assert 'base_url = "http://127.0.0.1:1111/v1"' not in content + + +def test_apply_codex_provider_scope_creates_new_config_when_missing( + monkeypatch, tmp_path: Path +) -> None: + config_path = tmp_path / "nested" / "config.toml" + monkeypatch.setattr("headroom.providers.codex.install.codex_config_path", lambda: config_path) + manifest = _manifest(tmp_path) + + mutation = apply_codex_provider_scope(manifest) + + assert mutation is not None + assert 'base_url = "http://127.0.0.1:8787/v1"' in config_path.read_text() + + +def test_revert_codex_provider_scope_ignores_missing_path_and_file(tmp_path: Path) -> None: + manifest = _manifest(tmp_path) + + revert_codex_provider_scope( + ManagedMutation(target="codex", kind="toml-block"), + manifest, + ) + revert_codex_provider_scope( + ManagedMutation( + target="codex", + kind="toml-block", + path=str(tmp_path / "missing.toml"), + ), + manifest, + ) + + +def test_revert_codex_provider_scope_ignores_files_without_managed_block( + monkeypatch, tmp_path: Path +) -> None: + config_path = tmp_path / "config.toml" + config_path.write_text('model = "gpt-4o"\n') + monkeypatch.setattr("headroom.providers.codex.install.codex_config_path", lambda: config_path) + manifest = _manifest(tmp_path) + mutation = ManagedMutation(target="codex", kind="toml-block", path=str(config_path)) + + revert_codex_provider_scope(mutation, manifest) + + assert config_path.read_text() == 'model = "gpt-4o"\n' + + +def test_apply_openclaw_provider_scope_uses_manifest_port(monkeypatch, tmp_path: Path) -> None: + recorded: list[list[str]] = [] + monkeypatch.setattr("headroom.providers.openclaw.install.shutil_which", lambda name: "openclaw") + monkeypatch.setattr( + "headroom.providers.openclaw.install.resolve_headroom_command", + lambda: ["headroom"], + ) + monkeypatch.setattr( + "headroom.providers.openclaw.install._invoke_openclaw", + lambda command: recorded.append(command), + ) + monkeypatch.setattr( + "headroom.providers.openclaw.install.openclaw_config_path", + lambda: tmp_path / "openclaw.json", + ) + manifest = _manifest(tmp_path) + manifest.port = 9999 + + from headroom.providers.openclaw.install import ( + apply_provider_scope as apply_openclaw_provider_scope, + ) + + apply_openclaw_provider_scope(manifest) + + assert recorded == [["headroom", "wrap", "openclaw", "--no-auto-start", "--proxy-port", "9999"]] + + +def test_openclaw_apply_provider_scope_requires_installed_binary( + tmp_path: Path, monkeypatch +) -> None: + monkeypatch.setattr("headroom.providers.openclaw.install.shutil_which", lambda name: None) + + with pytest.raises(click.ClickException, match="openclaw not found"): + from headroom.providers.openclaw.install import ( + apply_provider_scope as apply_openclaw_provider_scope, + ) + + apply_openclaw_provider_scope(_manifest(tmp_path)) + + +def test_openclaw_helper_wrappers_delegate_to_stdlib(monkeypatch) -> None: + monkeypatch.setattr("shutil.which", lambda name: f"/fake/{name}") + recorded: list[tuple[list[str], bool]] = [] + + def fake_run(command: list[str], check: bool) -> None: + recorded.append((command, check)) + + monkeypatch.setattr("subprocess.run", fake_run) + + from headroom.providers.openclaw.install import _invoke_openclaw, shutil_which + + assert shutil_which("openclaw") == "/fake/openclaw" + _invoke_openclaw(["headroom", "wrap", "openclaw"]) + + assert recorded == [(["headroom", "wrap", "openclaw"], True)] + + +def test_openclaw_revert_provider_scope_skips_without_binary(monkeypatch, tmp_path: Path) -> None: + monkeypatch.setattr("headroom.providers.openclaw.install.shutil_which", lambda name: None) + called = False + + def fail_if_called(command: list[str]) -> None: + nonlocal called + called = True + + monkeypatch.setattr("headroom.providers.openclaw.install._invoke_openclaw", fail_if_called) + + from headroom.providers.openclaw.install import ( + revert_provider_scope as revert_openclaw_provider_scope, + ) + + revert_openclaw_provider_scope( + ManagedMutation(target="openclaw", kind="openclaw-wrap", path=str(tmp_path / "cfg.json")), + _manifest(tmp_path), + ) + + assert called is False + + +def test_openclaw_revert_provider_scope_invokes_unwrap(monkeypatch, tmp_path: Path) -> None: + recorded: list[list[str]] = [] + monkeypatch.setattr("headroom.providers.openclaw.install.shutil_which", lambda name: "openclaw") + monkeypatch.setattr( + "headroom.providers.openclaw.install.resolve_headroom_command", + lambda: ["headroom"], + ) + monkeypatch.setattr( + "headroom.providers.openclaw.install._invoke_openclaw", + lambda command: recorded.append(command), + ) + + from headroom.providers.openclaw.install import ( + revert_provider_scope as revert_openclaw_provider_scope, + ) + + revert_openclaw_provider_scope( + ManagedMutation(target="openclaw", kind="openclaw-wrap", path=str(tmp_path / "cfg.json")), + _manifest(tmp_path), + ) + + assert recorded == [["headroom", "unwrap", "openclaw"]] + + +def test_windows_env_scope_restores_previous_values(monkeypatch, tmp_path: Path) -> None: + manifest = _manifest(tmp_path) + manifest.scope = "user" + manifest.targets = ["claude"] + manifest.base_env = {"HEADROOM_PORT": "8787"} + manifest.tool_envs = {"claude": {"ANTHROPIC_BASE_URL": "http://127.0.0.1:8787"}} + + calls: list[list[str]] = [] + previous_values = { + "HEADROOM_PORT": "7777", + "ANTHROPIC_BASE_URL": "https://old", + } + + class Result: + def __init__(self, stdout: str = "") -> None: + self.stdout = stdout + + def fake_run(command: list[str], **kwargs): + calls.append(command) + script = command[-1] + if "GetEnvironmentVariable" in script: + name = script.split("GetEnvironmentVariable('", 1)[1].split("'", 1)[0] + value = previous_values.get(name, "__HEADROOM_UNSET__") + return Result(stdout=value) + return Result() + + monkeypatch.setattr("headroom.install.providers.subprocess.run", fake_run) + + mutations = _apply_windows_env_scope(manifest) + _remove_windows_env_scope(mutations) + + previous_by_name = {mutation.data["name"]: mutation.data["previous"] for mutation in mutations} + assert previous_by_name["HEADROOM_PORT"] == "7777" + assert previous_by_name["ANTHROPIC_BASE_URL"] == "https://old" + assert any( + "[Environment]::SetEnvironmentVariable('HEADROOM_PORT','7777','User')" in command[-1] + for command in calls + ) + assert any( + "[Environment]::SetEnvironmentVariable('ANTHROPIC_BASE_URL','https://old','User')" + in command[-1] + for command in calls + ) + + +def test_remove_windows_env_scope_requires_name_and_scope() -> None: + try: + _remove_windows_env_scope([ManagedMutation(target="env", kind="windows-env", data={})]) + except ValueError as exc: + assert "variable name" in str(exc) + else: + raise AssertionError("expected missing variable name to raise") + + try: + _remove_windows_env_scope( + [ManagedMutation(target="env", kind="windows-env", data={"name": "X", "scope": 1})] + ) + except ValueError as exc: + assert "valid scope" in str(exc) + else: + raise AssertionError("expected invalid scope to raise") + + +def test_apply_mutations_runs_openclaw_for_user_scope(monkeypatch, tmp_path: Path) -> None: + manifest = _manifest(tmp_path) + manifest.scope = "user" + manifest.targets = ["openclaw"] + manifest.base_env = {"HEADROOM_PORT": "8787"} + manifest.tool_envs = {} + + if os.name == "nt": + monkeypatch.setattr( + "headroom.install.providers._apply_windows_env_scope", lambda deployment: [] + ) + else: + monkeypatch.setattr( + "headroom.install.providers._apply_unix_env_scope", lambda deployment: [] + ) + monkeypatch.setattr( + "headroom.install.providers.apply_provider_scope_mutations", + lambda deployment: [ManagedMutation(target="openclaw", kind="openclaw-wrap")], + ) + + from headroom.install.providers import apply_mutations + + mutations = apply_mutations(manifest) + + assert [mutation.kind for mutation in mutations] == ["openclaw-wrap"] + + +def test_claude_build_install_env_returns_proxy_base_url() -> None: + # Arrange / Act + env = build_claude_install_env(port=5566, backend="ignored") + + # Assert + assert env == {"ANTHROPIC_BASE_URL": "http://127.0.0.1:5566"} + + +def test_copilot_build_install_env_uses_provider_type_specific_proxy_urls() -> None: + anthropic_env = build_copilot_install_env(port=8787, backend="anthropic") + openai_env = build_copilot_install_env(port=8787, backend="anyllm") + + assert anthropic_env == { + "COPILOT_PROVIDER_TYPE": "anthropic", + "COPILOT_PROVIDER_BASE_URL": "http://127.0.0.1:8787", + } + assert openai_env == { + "COPILOT_PROVIDER_TYPE": "openai", + "COPILOT_PROVIDER_BASE_URL": "http://127.0.0.1:8787/v1", + "COPILOT_PROVIDER_WIRE_API": "completions", + } + + +def test_apply_claude_provider_scope_skips_non_provider_scope(monkeypatch, tmp_path: Path) -> None: + # Arrange + settings_path = tmp_path / "settings.json" + monkeypatch.setattr( + "headroom.providers.claude.install.claude_settings_path", lambda: settings_path + ) + manifest = _manifest(tmp_path) + manifest.scope = "user" + + # Act + mutation = apply_claude_provider_scope(manifest) + + # Assert + assert mutation is None + assert not settings_path.exists() + + +def test_revert_claude_provider_scope_removes_new_values_from_non_mapping_env( + monkeypatch, tmp_path: Path +) -> None: + # Arrange + settings_path = tmp_path / "settings.json" + settings_path.write_text(json.dumps({"env": ["not-a-map"]})) + monkeypatch.setattr( + "headroom.providers.claude.install.claude_settings_path", lambda: settings_path + ) + manifest = _manifest(tmp_path) + + # Act + mutation = apply_claude_provider_scope(manifest) + apply_payload = json.loads(settings_path.read_text()) + revert_claude_provider_scope(mutation, manifest) + reverted_payload = json.loads(settings_path.read_text()) + + # Assert + assert mutation is not None + assert mutation.data["previous"] == {"ANTHROPIC_BASE_URL": None} + assert apply_payload["env"] == {"ANTHROPIC_BASE_URL": "http://127.0.0.1:8787"} + assert reverted_payload["env"] == {} + + +def test_apply_claude_provider_scope_creates_settings_when_missing( + monkeypatch, tmp_path: Path +) -> None: + # Arrange + settings_path = tmp_path / "nested" / "settings.json" + monkeypatch.setattr( + "headroom.providers.claude.install.claude_settings_path", lambda: settings_path + ) + manifest = _manifest(tmp_path) + + # Act + mutation = apply_claude_provider_scope(manifest) + + # Assert + assert mutation is not None + assert json.loads(settings_path.read_text()) == { + "env": {"ANTHROPIC_BASE_URL": "http://127.0.0.1:8787"} + } + + +def test_revert_claude_provider_scope_ignores_missing_mutation_path(tmp_path: Path) -> None: + # Arrange + manifest = _manifest(tmp_path) + mutation = ManagedMutation(target="claude", kind="json-env", data={"previous": {}}) + + # Act / Assert + revert_claude_provider_scope(mutation, manifest) + + +def test_revert_claude_provider_scope_ignores_missing_settings_file(tmp_path: Path) -> None: + # Arrange + manifest = _manifest(tmp_path) + mutation = ManagedMutation( + target="claude", + kind="json-env", + path=str(tmp_path / "missing-settings.json"), + data={"previous": {}}, + ) + + # Act / Assert + revert_claude_provider_scope(mutation, manifest) diff --git a/tests/test_provider_cursor.py b/tests/test_provider_cursor.py index 1f8a71b8b..b16671524 100644 --- a/tests/test_provider_cursor.py +++ b/tests/test_provider_cursor.py @@ -1,30 +1,30 @@ -from __future__ import annotations - -from headroom.providers.cursor import build_proxy_targets, render_setup_lines -from headroom.providers.cursor.install import build_install_env - - -def test_cursor_proxy_targets_use_local_headroom_proxy() -> None: - targets = build_proxy_targets(9999) - - assert targets.openai_base_url == "http://127.0.0.1:9999/v1" - assert targets.anthropic_base_url == "http://127.0.0.1:9999" - - -def test_cursor_setup_lines_include_both_provider_urls() -> None: - lines = render_setup_lines(8787) - joined = "\n".join(lines) - - assert "http://127.0.0.1:8787/v1" in joined - assert "http://127.0.0.1:8787" in joined - - -def test_cursor_build_install_env_returns_both_proxy_urls() -> None: - # Arrange / Act - env = build_install_env(port=7654, backend="ignored") - - # Assert - assert env == { - "OPENAI_BASE_URL": "http://127.0.0.1:7654/v1", - "ANTHROPIC_BASE_URL": "http://127.0.0.1:7654", - } +from __future__ import annotations + +from headroom.providers.cursor import build_proxy_targets, render_setup_lines +from headroom.providers.cursor.install import build_install_env + + +def test_cursor_proxy_targets_use_local_headroom_proxy() -> None: + targets = build_proxy_targets(9999) + + assert targets.openai_base_url == "http://127.0.0.1:9999/v1" + assert targets.anthropic_base_url == "http://127.0.0.1:9999" + + +def test_cursor_setup_lines_include_both_provider_urls() -> None: + lines = render_setup_lines(8787) + joined = "\n".join(lines) + + assert "http://127.0.0.1:8787/v1" in joined + assert "http://127.0.0.1:8787" in joined + + +def test_cursor_build_install_env_returns_both_proxy_urls() -> None: + # Arrange / Act + env = build_install_env(port=7654, backend="ignored") + + # Assert + assert env == { + "OPENAI_BASE_URL": "http://127.0.0.1:7654/v1", + "ANTHROPIC_BASE_URL": "http://127.0.0.1:7654", + } diff --git a/tests/test_provider_registry.py b/tests/test_provider_registry.py index 9312b1b88..c9d782cd4 100644 --- a/tests/test_provider_registry.py +++ b/tests/test_provider_registry.py @@ -1,393 +1,393 @@ -from __future__ import annotations - -import logging - -from headroom.providers.registry import ( - ProviderApiOverrides, - build_proxy_provider_runtime, - create_proxy_backend, - format_backend_status, - resolve_api_overrides, - resolve_api_targets, -) -from headroom.proxy.models import ProxyConfig - - -def test_resolve_api_overrides_prefers_explicit_values_over_environment(monkeypatch) -> None: - monkeypatch.setenv("ANTHROPIC_TARGET_API_URL", "https://env.anthropic.example/v1") - monkeypatch.setenv("OPENAI_TARGET_API_URL", "https://env.openai.example/v1") - - overrides = resolve_api_overrides( - anthropic_api_url="https://cli.anthropic.example/v1", - openai_api_url=None, - gemini_api_url=None, - cloudcode_api_url=None, - ) - - assert overrides == ProviderApiOverrides( - anthropic="https://cli.anthropic.example/v1", - openai="https://env.openai.example/v1", - gemini=None, - cloudcode=None, - ) - - -def test_resolve_api_targets_normalizes_trailing_v1() -> None: - targets = resolve_api_targets( - ProviderApiOverrides( - anthropic="https://anthropic.example/v1/", - openai="https://openai.example/v1", - gemini="https://gemini.example/v1", - cloudcode="https://cloudcode.example/v1/", - ) - ) - - assert targets.anthropic == "https://anthropic.example" - assert targets.openai == "https://openai.example" - assert targets.gemini == "https://gemini.example" - assert targets.cloudcode == "https://cloudcode.example" - - -def test_proxy_config_exposes_provider_api_overrides() -> None: - config = ProxyConfig( - anthropic_api_url="https://anthropic.example", - openai_api_url="https://openai.example", - gemini_api_url=None, - cloudcode_api_url="https://cloudcode.example", - ) - - assert config.provider_api_overrides == ProviderApiOverrides( - anthropic="https://anthropic.example", - openai="https://openai.example", - gemini=None, - cloudcode="https://cloudcode.example", - ) - - -def test_format_backend_status_for_anyllm() -> None: - assert ( - format_backend_status( - backend="anyllm", - anyllm_provider="groq", - bedrock_region="us-central1", - ) - == "Groq via any-llm" - ) - - -def test_format_backend_status_for_anthropic_direct() -> None: - assert ( - format_backend_status( - backend="anthropic", - anyllm_provider="ignored", - bedrock_region=None, - ) - == "ANTHROPIC (direct API)" - ) - - -def test_proxy_provider_runtime_routes_model_metadata_and_passthrough() -> None: - runtime = build_proxy_provider_runtime(ProxyConfig()) - - assert runtime.model_metadata_provider({"x-api-key": "test"}) == "anthropic" - assert runtime.model_metadata_provider({}) == "openai" - assert ( - runtime.select_passthrough_base_url({"x-api-key": "test"}) == runtime.api_targets.anthropic - ) - assert ( - runtime.select_passthrough_base_url({"x-goog-api-key": "test"}) - == runtime.api_targets.gemini - ) - assert runtime.select_passthrough_base_url({"api-key": "azure", "x-headroom-base-url": ""}) == ( - runtime.api_targets.openai - ) - - -def test_create_proxy_backend_handles_missing_litellm_backend(caplog) -> None: - logger = logging.getLogger("test") - - with caplog.at_level(logging.WARNING): - missing = create_proxy_backend( - backend="bedrock", - anyllm_provider="ignored", - bedrock_region="us-east-1", - logger=logger, - litellm_backend_cls=lambda provider, region: (_ for _ in ()).throw( - ImportError("missing") - ), - ) - - assert missing is None - assert "LiteLLM backend not available" in caplog.text - - -def test_proxy_provider_runtime_loaders_cache_backend_types(monkeypatch) -> None: - import headroom.providers.registry as registry - - anyllm_loads = 0 - litellm_loads = 0 - - class FakeAnyLLMBackend: - pass - - class FakeLiteLLMBackend: - pass - - def fake_import(name, globals=None, locals=None, fromlist=(), level=0): - nonlocal anyllm_loads, litellm_loads - if name == "headroom.backends.anyllm": - anyllm_loads += 1 - return type("Module", (), {"AnyLLMBackend": FakeAnyLLMBackend})() - if name == "headroom.backends.litellm": - litellm_loads += 1 - return type("Module", (), {"LiteLLMBackend": FakeLiteLLMBackend})() - raise AssertionError(name) - - monkeypatch.setattr(registry, "AnyLLMBackendType", None) - monkeypatch.setattr(registry, "LiteLLMBackendType", None) - monkeypatch.setattr("builtins.__import__", fake_import) - - assert registry._load_anyllm_backend() is FakeAnyLLMBackend - assert registry._load_anyllm_backend() is FakeAnyLLMBackend - assert registry._load_litellm_backend() is FakeLiteLLMBackend - assert registry._load_litellm_backend() is FakeLiteLLMBackend - assert anyllm_loads == 1 - assert litellm_loads == 1 - - -def test_proxy_provider_runtime_transport_helpers_handle_missing_usage() -> None: - import headroom.providers.registry as registry - - class Storage: - def __init__(self) -> None: - self.saved = [] - - def save(self, metrics) -> None: - self.saved.append(metrics) - - client = type( - "Client", - (), - { - "_storage": Storage(), - "_original": type( - "Original", - (), - { - "chat": type( - "Chat", - (), - { - "completions": type( - "Completions", - (), - { - "create": staticmethod( - lambda **kwargs: type("Resp", (), {"usage": None})() - ) - }, - )() - }, - )(), - "messages": type( - "Messages", - (), - { - "create": staticmethod( - lambda **kwargs: type("Resp", (), {"usage": None})() - ) - }, - )(), - }, - )(), - }, - )() - openai_metrics = type("Metrics", (), {"tokens_output": 0, "cached_tokens": 0})() - anthropic_metrics = type("Metrics", (), {"tokens_output": 0, "cached_tokens": 0})() - - registry._call_openai_transport( - client, - model="gpt-4o", - messages=[], - stream=False, - metrics=openai_metrics, - ) - registry._call_anthropic_transport( - client, - model="claude", - messages=[], - stream=False, - metrics=anthropic_metrics, - ) - - assert openai_metrics.tokens_output == 0 - assert openai_metrics.cached_tokens == 0 - assert anthropic_metrics.tokens_output == 0 - assert anthropic_metrics.cached_tokens == 0 - assert len(client._storage.saved) == 2 - - -def test_proxy_provider_runtime_transport_helpers_handle_usage_without_optional_cache_fields() -> ( - None -): - import headroom.providers.registry as registry - - class Storage: - def __init__(self) -> None: - self.saved = [] - - def save(self, metrics) -> None: - self.saved.append(metrics) - - client = type( - "Client", - (), - { - "_storage": Storage(), - "_original": type( - "Original", - (), - { - "chat": type( - "Chat", - (), - { - "completions": type( - "Completions", - (), - { - "create": staticmethod( - lambda **kwargs: type( - "Resp", - (), - { - "usage": type( - "Usage", - (), - {"completion_tokens": 7}, - )() - }, - )() - ) - }, - )() - }, - )(), - "messages": type( - "Messages", - (), - { - "create": staticmethod( - lambda **kwargs: type( - "Resp", - (), - { - "usage": type( - "Usage", - (), - {"output_tokens": 5}, - )() - }, - )() - ) - }, - )(), - }, - )(), - }, - )() - openai_metrics = type("Metrics", (), {"tokens_output": 0, "cached_tokens": 0})() - anthropic_metrics = type("Metrics", (), {"tokens_output": 0, "cached_tokens": 0})() - - registry._call_openai_transport( - client, - model="gpt-4o", - messages=[], - stream=False, - metrics=openai_metrics, - ) - registry._call_anthropic_transport( - client, - model="claude", - messages=[], - stream=False, - metrics=anthropic_metrics, - ) - - assert openai_metrics.tokens_output == 7 - assert openai_metrics.cached_tokens == 0 - assert anthropic_metrics.tokens_output == 5 - assert anthropic_metrics.cached_tokens == 0 - assert len(client._storage.saved) == 2 - - -def test_proxy_provider_runtime_openai_transport_handles_prompt_details_without_cached_tokens() -> ( - None -): - import headroom.providers.registry as registry - - class Storage: - def __init__(self) -> None: - self.saved = [] - - def save(self, metrics) -> None: - self.saved.append(metrics) - - client = type( - "Client", - (), - { - "_storage": Storage(), - "_original": type( - "Original", - (), - { - "chat": type( - "Chat", - (), - { - "completions": type( - "Completions", - (), - { - "create": staticmethod( - lambda **kwargs: type( - "Resp", - (), - { - "usage": type( - "Usage", - (), - { - "completion_tokens": 9, - "prompt_tokens_details": type( - "Details", - (), - {}, - )(), - }, - )() - }, - )() - ) - }, - )() - }, - )() - }, - )(), - }, - )() - metrics = type("Metrics", (), {"tokens_output": 0, "cached_tokens": 0})() - - registry._call_openai_transport( - client, - model="gpt-4o", - messages=[], - stream=False, - metrics=metrics, - ) - - assert metrics.tokens_output == 9 - assert metrics.cached_tokens == 0 - assert len(client._storage.saved) == 1 +from __future__ import annotations + +import logging + +from headroom.providers.registry import ( + ProviderApiOverrides, + build_proxy_provider_runtime, + create_proxy_backend, + format_backend_status, + resolve_api_overrides, + resolve_api_targets, +) +from headroom.proxy.models import ProxyConfig + + +def test_resolve_api_overrides_prefers_explicit_values_over_environment(monkeypatch) -> None: + monkeypatch.setenv("ANTHROPIC_TARGET_API_URL", "https://env.anthropic.example/v1") + monkeypatch.setenv("OPENAI_TARGET_API_URL", "https://env.openai.example/v1") + + overrides = resolve_api_overrides( + anthropic_api_url="https://cli.anthropic.example/v1", + openai_api_url=None, + gemini_api_url=None, + cloudcode_api_url=None, + ) + + assert overrides == ProviderApiOverrides( + anthropic="https://cli.anthropic.example/v1", + openai="https://env.openai.example/v1", + gemini=None, + cloudcode=None, + ) + + +def test_resolve_api_targets_normalizes_trailing_v1() -> None: + targets = resolve_api_targets( + ProviderApiOverrides( + anthropic="https://anthropic.example/v1/", + openai="https://openai.example/v1", + gemini="https://gemini.example/v1", + cloudcode="https://cloudcode.example/v1/", + ) + ) + + assert targets.anthropic == "https://anthropic.example" + assert targets.openai == "https://openai.example" + assert targets.gemini == "https://gemini.example" + assert targets.cloudcode == "https://cloudcode.example" + + +def test_proxy_config_exposes_provider_api_overrides() -> None: + config = ProxyConfig( + anthropic_api_url="https://anthropic.example", + openai_api_url="https://openai.example", + gemini_api_url=None, + cloudcode_api_url="https://cloudcode.example", + ) + + assert config.provider_api_overrides == ProviderApiOverrides( + anthropic="https://anthropic.example", + openai="https://openai.example", + gemini=None, + cloudcode="https://cloudcode.example", + ) + + +def test_format_backend_status_for_anyllm() -> None: + assert ( + format_backend_status( + backend="anyllm", + anyllm_provider="groq", + bedrock_region="us-central1", + ) + == "Groq via any-llm" + ) + + +def test_format_backend_status_for_anthropic_direct() -> None: + assert ( + format_backend_status( + backend="anthropic", + anyllm_provider="ignored", + bedrock_region=None, + ) + == "ANTHROPIC (direct API)" + ) + + +def test_proxy_provider_runtime_routes_model_metadata_and_passthrough() -> None: + runtime = build_proxy_provider_runtime(ProxyConfig()) + + assert runtime.model_metadata_provider({"x-api-key": "test"}) == "anthropic" + assert runtime.model_metadata_provider({}) == "openai" + assert ( + runtime.select_passthrough_base_url({"x-api-key": "test"}) == runtime.api_targets.anthropic + ) + assert ( + runtime.select_passthrough_base_url({"x-goog-api-key": "test"}) + == runtime.api_targets.gemini + ) + assert runtime.select_passthrough_base_url({"api-key": "azure", "x-headroom-base-url": ""}) == ( + runtime.api_targets.openai + ) + + +def test_create_proxy_backend_handles_missing_litellm_backend(caplog) -> None: + logger = logging.getLogger("test") + + with caplog.at_level(logging.WARNING): + missing = create_proxy_backend( + backend="bedrock", + anyllm_provider="ignored", + bedrock_region="us-east-1", + logger=logger, + litellm_backend_cls=lambda provider, region: (_ for _ in ()).throw( + ImportError("missing") + ), + ) + + assert missing is None + assert "LiteLLM backend not available" in caplog.text + + +def test_proxy_provider_runtime_loaders_cache_backend_types(monkeypatch) -> None: + import headroom.providers.registry as registry + + anyllm_loads = 0 + litellm_loads = 0 + + class FakeAnyLLMBackend: + pass + + class FakeLiteLLMBackend: + pass + + def fake_import(name, globals=None, locals=None, fromlist=(), level=0): + nonlocal anyllm_loads, litellm_loads + if name == "headroom.backends.anyllm": + anyllm_loads += 1 + return type("Module", (), {"AnyLLMBackend": FakeAnyLLMBackend})() + if name == "headroom.backends.litellm": + litellm_loads += 1 + return type("Module", (), {"LiteLLMBackend": FakeLiteLLMBackend})() + raise AssertionError(name) + + monkeypatch.setattr(registry, "AnyLLMBackendType", None) + monkeypatch.setattr(registry, "LiteLLMBackendType", None) + monkeypatch.setattr("builtins.__import__", fake_import) + + assert registry._load_anyllm_backend() is FakeAnyLLMBackend + assert registry._load_anyllm_backend() is FakeAnyLLMBackend + assert registry._load_litellm_backend() is FakeLiteLLMBackend + assert registry._load_litellm_backend() is FakeLiteLLMBackend + assert anyllm_loads == 1 + assert litellm_loads == 1 + + +def test_proxy_provider_runtime_transport_helpers_handle_missing_usage() -> None: + import headroom.providers.registry as registry + + class Storage: + def __init__(self) -> None: + self.saved = [] + + def save(self, metrics) -> None: + self.saved.append(metrics) + + client = type( + "Client", + (), + { + "_storage": Storage(), + "_original": type( + "Original", + (), + { + "chat": type( + "Chat", + (), + { + "completions": type( + "Completions", + (), + { + "create": staticmethod( + lambda **kwargs: type("Resp", (), {"usage": None})() + ) + }, + )() + }, + )(), + "messages": type( + "Messages", + (), + { + "create": staticmethod( + lambda **kwargs: type("Resp", (), {"usage": None})() + ) + }, + )(), + }, + )(), + }, + )() + openai_metrics = type("Metrics", (), {"tokens_output": 0, "cached_tokens": 0})() + anthropic_metrics = type("Metrics", (), {"tokens_output": 0, "cached_tokens": 0})() + + registry._call_openai_transport( + client, + model="gpt-4o", + messages=[], + stream=False, + metrics=openai_metrics, + ) + registry._call_anthropic_transport( + client, + model="claude", + messages=[], + stream=False, + metrics=anthropic_metrics, + ) + + assert openai_metrics.tokens_output == 0 + assert openai_metrics.cached_tokens == 0 + assert anthropic_metrics.tokens_output == 0 + assert anthropic_metrics.cached_tokens == 0 + assert len(client._storage.saved) == 2 + + +def test_proxy_provider_runtime_transport_helpers_handle_usage_without_optional_cache_fields() -> ( + None +): + import headroom.providers.registry as registry + + class Storage: + def __init__(self) -> None: + self.saved = [] + + def save(self, metrics) -> None: + self.saved.append(metrics) + + client = type( + "Client", + (), + { + "_storage": Storage(), + "_original": type( + "Original", + (), + { + "chat": type( + "Chat", + (), + { + "completions": type( + "Completions", + (), + { + "create": staticmethod( + lambda **kwargs: type( + "Resp", + (), + { + "usage": type( + "Usage", + (), + {"completion_tokens": 7}, + )() + }, + )() + ) + }, + )() + }, + )(), + "messages": type( + "Messages", + (), + { + "create": staticmethod( + lambda **kwargs: type( + "Resp", + (), + { + "usage": type( + "Usage", + (), + {"output_tokens": 5}, + )() + }, + )() + ) + }, + )(), + }, + )(), + }, + )() + openai_metrics = type("Metrics", (), {"tokens_output": 0, "cached_tokens": 0})() + anthropic_metrics = type("Metrics", (), {"tokens_output": 0, "cached_tokens": 0})() + + registry._call_openai_transport( + client, + model="gpt-4o", + messages=[], + stream=False, + metrics=openai_metrics, + ) + registry._call_anthropic_transport( + client, + model="claude", + messages=[], + stream=False, + metrics=anthropic_metrics, + ) + + assert openai_metrics.tokens_output == 7 + assert openai_metrics.cached_tokens == 0 + assert anthropic_metrics.tokens_output == 5 + assert anthropic_metrics.cached_tokens == 0 + assert len(client._storage.saved) == 2 + + +def test_proxy_provider_runtime_openai_transport_handles_prompt_details_without_cached_tokens() -> ( + None +): + import headroom.providers.registry as registry + + class Storage: + def __init__(self) -> None: + self.saved = [] + + def save(self, metrics) -> None: + self.saved.append(metrics) + + client = type( + "Client", + (), + { + "_storage": Storage(), + "_original": type( + "Original", + (), + { + "chat": type( + "Chat", + (), + { + "completions": type( + "Completions", + (), + { + "create": staticmethod( + lambda **kwargs: type( + "Resp", + (), + { + "usage": type( + "Usage", + (), + { + "completion_tokens": 9, + "prompt_tokens_details": type( + "Details", + (), + {}, + )(), + }, + )() + }, + )() + ) + }, + )() + }, + )() + }, + )(), + }, + )() + metrics = type("Metrics", (), {"tokens_output": 0, "cached_tokens": 0})() + + registry._call_openai_transport( + client, + model="gpt-4o", + messages=[], + stream=False, + metrics=metrics, + ) + + assert metrics.tokens_output == 9 + assert metrics.cached_tokens == 0 + assert len(client._storage.saved) == 1 diff --git a/tests/test_providers/test_universal.py b/tests/test_providers/test_universal.py index cbfb9c15a..1b83a678c 100644 --- a/tests/test_providers/test_universal.py +++ b/tests/test_providers/test_universal.py @@ -1,495 +1,495 @@ -"""Tests for universal provider support. - -Tests OpenAICompatibleProvider, GoogleProvider, and LiteLLMProvider. -""" - -from __future__ import annotations - -import pytest - -from headroom.providers import ( - GoogleProvider, - LiteLLMProvider, - ModelCapabilities, - OpenAICompatibleProvider, - create_anyscale_provider, - create_fireworks_provider, - create_groq_provider, - create_litellm_provider, - create_lmstudio_provider, - create_ollama_provider, - create_together_provider, - create_vllm_provider, - is_litellm_available, -) - - -def _transformers_available() -> bool: - """Check if transformers is available.""" - try: - import transformers # noqa: F401 - - return True - except ImportError: - return False - - -class TestOpenAICompatibleProvider: - """Tests for OpenAICompatibleProvider.""" - - def test_init_default(self): - """Test initialization with defaults.""" - provider = OpenAICompatibleProvider() - assert provider.name == "openai_compatible" - assert provider.base_url is None - - def test_init_with_config(self): - """Test initialization with configuration.""" - provider = OpenAICompatibleProvider( - name="custom", - base_url="http://localhost:8080/v1", - api_key="test-key", - ) - assert provider.name == "custom" - assert provider.base_url == "http://localhost:8080/v1" - assert provider.api_key == "test-key" - - def test_supports_any_model(self): - """Test that provider supports any model.""" - provider = OpenAICompatibleProvider() - assert provider.supports_model("any-model") is True - assert provider.supports_model("llama-3") is True - assert provider.supports_model("custom-finetuned") is True - - @pytest.mark.skipif( - not _transformers_available(), - reason="transformers not installed - needed for HuggingFace tokenizer", - ) - def test_get_token_counter(self): - """Test getting token counter.""" - provider = OpenAICompatibleProvider() - counter = provider.get_token_counter("llama-3-8b") - assert counter is not None - # Should be able to count tokens - count = counter.count_text("Hello, world!") - assert count > 0 - - def test_get_context_limit_known_model(self): - """Test context limit for known models.""" - provider = OpenAICompatibleProvider() - # Llama 3.1 has 128K context - limit = provider.get_context_limit("llama-3.1-8b") - assert limit == 128000 - - def test_get_context_limit_unknown_model(self): - """Test context limit for unknown models (defaults to 128K).""" - provider = OpenAICompatibleProvider() - limit = provider.get_context_limit("unknown-model") - assert limit == 128000 - - def test_register_model(self): - """Test registering a custom model.""" - provider = OpenAICompatibleProvider() - provider.register_model( - "my-model", - context_window=64000, - max_output_tokens=8192, - input_cost_per_1m=1.0, - output_cost_per_1m=2.0, - ) - assert provider.get_context_limit("my-model") == 64000 - - def test_estimate_cost_registered_model(self): - """Test cost estimation for registered model.""" - provider = OpenAICompatibleProvider() - provider.register_model( - "priced-model", - input_cost_per_1m=1.0, - output_cost_per_1m=2.0, - ) - cost = provider.estimate_cost( - input_tokens=1000000, - output_tokens=500000, - model="priced-model", - ) - assert cost == 2.0 # 1.0 + 1.0 - - def test_estimate_cost_unknown_model(self): - """Test cost estimation returns None for unknown model.""" - provider = OpenAICompatibleProvider() - cost = provider.estimate_cost( - input_tokens=1000, - output_tokens=500, - model="unknown-model", - ) - assert cost is None - - def test_register_model_accepts_capabilities_object(self): - provider = OpenAICompatibleProvider() - caps = ModelCapabilities(model="caps-model", context_window=16000, tokenizer_backend="test") - - provider.register_model("caps-model", capabilities=caps) - - assert provider.get_context_limit("caps-model") == 16000 - - def test_get_token_counter_uses_registered_tokenizer_backend(self, monkeypatch): - recorded: list[tuple[str, str | None]] = [] - - class DummyTokenizer: - def count_text(self, text: str) -> int: - return len(text.split()) - - monkeypatch.setattr( - "headroom.providers.openai_compatible.get_tokenizer", - lambda model, backend=None: recorded.append((model, backend)) or DummyTokenizer(), - ) - provider = OpenAICompatibleProvider( - models={ - "custom-model": ModelCapabilities( - model="custom-model", - tokenizer_backend="custom-backend", - ) - } - ) - - counter = provider.get_token_counter("custom-model") - - assert counter.count_text("one two three") == 3 - assert recorded == [("custom-model", "custom-backend")] - - def test_openai_compatible_token_counter_counts_message_parts(self, monkeypatch): - class DummyTokenizer: - def count_text(self, text: str) -> int: - return len(text) - - monkeypatch.setattr( - "headroom.providers.openai_compatible.get_tokenizer", - lambda model, backend=None: DummyTokenizer(), - ) - counter = OpenAICompatibleProvider().get_token_counter("demo-model") - - tokens = counter.count_message( - { - "role": "user", - "content": [{"type": "text", "text": "hi"}, "there"], - "name": "tester", - "tool_calls": [{"function": {"name": "lookup", "arguments": '{"x":1}'}}], - "tool_call_id": "call_123", - } - ) - total = counter.count_messages( - [ - {"role": "user", "content": "hello"}, - {"role": "assistant", "content": ["world"]}, - ] - ) - - assert tokens == 55 - assert total == 34 - - def test_openai_compatible_token_counter_ignores_unhandled_content_shapes(self, monkeypatch): - class DummyTokenizer: - def count_text(self, text: str) -> int: - return len(text) - - monkeypatch.setattr( - "headroom.providers.openai_compatible.get_tokenizer", - lambda model, backend=None: DummyTokenizer(), - ) - counter = OpenAICompatibleProvider().get_token_counter("demo-model") - - assert counter.count_message({"role": "user", "content": {}}) == 8 - assert counter.count_message({"role": "user", "content": [{"type": "image"}, 123]}) == 8 - - def test_get_context_limit_prefix_output_buffer_and_partial_pricing(self): - provider = OpenAICompatibleProvider( - models={ - "buffered": ModelCapabilities( - model="buffered", - max_output_tokens=1200, - input_cost_per_1m=1.0, - ) - } - ) - - assert provider.get_context_limit("mistral-custom") == 32768 - assert provider.get_output_buffer("buffered", default=4000) == 1200 - assert provider.get_output_buffer("unknown", default=2222) == 2222 - assert provider.estimate_cost(1000, 1000, "buffered") is None - - -class TestModelCapabilities: - """Tests for ModelCapabilities dataclass.""" - - def test_default_values(self): - """Test default capability values.""" - caps = ModelCapabilities(model="test-model") - assert caps.context_window == 128000 - assert caps.max_output_tokens == 4096 - assert caps.supports_tools is True - assert caps.supports_vision is False - assert caps.supports_streaming is True - - def test_custom_values(self): - """Test custom capability values.""" - caps = ModelCapabilities( - model="custom-model", - context_window=32000, - max_output_tokens=16384, - supports_tools=False, - supports_vision=True, - input_cost_per_1m=0.5, - output_cost_per_1m=1.5, - ) - assert caps.context_window == 32000 - assert caps.max_output_tokens == 16384 - assert caps.supports_tools is False - assert caps.supports_vision is True - assert caps.input_cost_per_1m == 0.5 - assert caps.output_cost_per_1m == 1.5 - - -class TestGoogleProvider: - """Tests for GoogleProvider.""" - - @pytest.fixture - def provider(self): - """Create Google provider.""" - return GoogleProvider() - - def test_name(self, provider): - """Test provider name.""" - assert provider.name == "google" - - def test_supports_gemini_models(self, provider): - """Test support for Gemini models.""" - assert provider.supports_model("gemini-2.0-flash") is True - assert provider.supports_model("gemini-1.5-pro") is True - assert provider.supports_model("gemini-1.5-flash") is True - - def test_not_supports_other_models(self, provider): - """Test non-support for other models.""" - assert provider.supports_model("gpt-4o") is False - assert provider.supports_model("claude-3") is False - - def test_get_token_counter(self, provider): - """Test getting token counter.""" - counter = provider.get_token_counter("gemini-2.0-flash") - assert counter is not None - count = counter.count_text("Hello, world!") - assert count > 0 - - def test_get_context_limit_gemini_2(self, provider): - """Test context limit for Gemini 2.0.""" - limit = provider.get_context_limit("gemini-2.0-flash") - # LiteLLM returns 1048576 (2^20), fallback returns 1000000 - assert limit in (1000000, 1048576) # ~1M tokens - - def test_get_context_limit_gemini_1_5_pro(self, provider): - """Test context limit for Gemini 1.5 Pro (2M!).""" - limit = provider.get_context_limit("gemini-1.5-pro") - # LiteLLM returns 2097152 (2^21), fallback returns 2000000 - assert limit in (2000000, 2097152) # ~2M tokens! - - def test_estimate_cost(self, provider): - """Test cost estimation.""" - cost = provider.estimate_cost( - input_tokens=1000000, - output_tokens=500000, - model="gemini-2.0-flash", - ) - assert cost is not None - # 1M input * $0.10 + 0.5M output * $0.40 = $0.10 + $0.20 = $0.30 - assert abs(cost - 0.30) < 0.01 - - def test_openai_compatible_url(self): - """Test OpenAI-compatible URL.""" - url = GoogleProvider.get_openai_compatible_url("test-key") - assert "generativelanguage.googleapis.com" in url - - -class TestProviderFactoryFunctions: - """Tests for provider factory functions.""" - - def test_create_ollama_provider(self): - """Test creating Ollama provider.""" - provider = create_ollama_provider() - assert provider.name == "ollama" - assert provider.base_url == "http://localhost:11434/v1" - - def test_create_ollama_provider_custom_url(self): - """Test creating Ollama provider with custom URL.""" - provider = create_ollama_provider("http://192.168.1.100:11434/v1") - assert provider.base_url == "http://192.168.1.100:11434/v1" - - def test_create_together_provider(self): - """Test creating Together provider.""" - provider = create_together_provider() - assert provider.name == "together" - assert "together.xyz" in provider.base_url - - def test_create_groq_provider(self): - """Test creating Groq provider.""" - provider = create_groq_provider() - assert provider.name == "groq" - assert "groq.com" in provider.base_url - - def test_create_vllm_provider(self): - """Test creating vLLM provider.""" - provider = create_vllm_provider("http://localhost:8000/v1") - assert provider.name == "vllm" - assert provider.base_url == "http://localhost:8000/v1" - - def test_create_lmstudio_provider(self): - """Test creating LM Studio provider.""" - provider = create_lmstudio_provider() - assert provider.name == "lmstudio" - assert provider.base_url == "http://localhost:1234/v1" - - def test_create_fireworks_and_anyscale_providers(self): - fireworks = create_fireworks_provider(api_key="fireworks-key") - anyscale = create_anyscale_provider(api_key="anyscale-key") - - assert fireworks.name == "fireworks" - assert fireworks.base_url == "https://api.fireworks.ai/inference/v1" - assert fireworks.api_key == "fireworks-key" - assert anyscale.name == "anyscale" - assert anyscale.base_url == "https://api.endpoints.anyscale.com/v1" - assert anyscale.api_key == "anyscale-key" - - -class TestLiteLLMProvider: - """Tests for LiteLLM provider.""" - - def test_is_litellm_available(self): - """Test checking LiteLLM availability.""" - result = is_litellm_available() - assert isinstance(result, bool) - - def test_unavailable_litellm_paths(self, monkeypatch): - import headroom.providers.litellm as litellm_module - - monkeypatch.setattr(litellm_module, "LITELLM_AVAILABLE", False) - - assert litellm_module.is_litellm_available() is False - assert litellm_module.LiteLLMProvider.list_supported_providers() == [] - with pytest.raises(RuntimeError, match="LiteLLM is required"): - litellm_module.LiteLLMTokenCounter("gpt-4o") - with pytest.raises(RuntimeError, match="LiteLLM is required"): - litellm_module.LiteLLMProvider() - - def test_litellm_token_counter_fallback_paths(self, monkeypatch): - import headroom.providers.litellm as litellm_module - - class DummyFallback: - def count_text(self, text: str) -> int: - return len(text.split()) - - monkeypatch.setattr(litellm_module, "LITELLM_AVAILABLE", True) - monkeypatch.setattr( - litellm_module, - "litellm_token_counter", - lambda **kwargs: (_ for _ in ()).throw(RuntimeError("boom")), - ) - monkeypatch.setattr(litellm_module, "EstimatingTokenCounter", DummyFallback) - - counter = litellm_module.LiteLLMTokenCounter("gpt-4o") - - assert counter.count_text("") == 0 - assert counter.count_text("one two three") == 3 - assert counter.count_message({"content": "one two"}) == 6 - assert counter.count_messages([]) == 0 - assert counter.count_messages([{"content": "one two"}, {"content": "three"}]) == 14 - - def test_litellm_provider_info_and_cost_fallbacks(self, monkeypatch): - import headroom.providers.litellm as litellm_module - - monkeypatch.setattr(litellm_module, "LITELLM_AVAILABLE", True) - monkeypatch.setattr( - litellm_module, - "litellm_get_model_info", - lambda model: { - "ctx-model": {"max_input_tokens": 64000}, - "max-model": {"max_tokens": 32000}, - "none-model": {"max_input_tokens": None, "max_output_tokens": None}, - "output-model": {"max_output_tokens": 6000}, - }[model], - ) - monkeypatch.setattr( - litellm_module, - "litellm", - type( - "LiteLLM", - (), - { - "completion_cost": staticmethod( - lambda **kwargs: 1.23 - if kwargs["model"] == "priced-model" - else (_ for _ in ()).throw(RuntimeError("missing price")) - ) - }, - )(), - ) - - provider = litellm_module.LiteLLMProvider() - - assert provider.get_context_limit("ctx-model") == 64000 - assert provider.get_context_limit("max-model") == 32000 - assert provider.get_context_limit("none-model") == 128000 - assert provider.get_output_buffer("output-model", default=4000) == 4000 - assert provider.get_output_buffer("none-model", default=2222) == 2222 - assert provider.estimate_cost(1000, 1000, "priced-model") == 1.23 - assert provider.estimate_cost(1000, 1000, "missing-price") is None - - def test_litellm_provider_handles_info_exceptions_and_factory(self, monkeypatch): - import headroom.providers.litellm as litellm_module - - monkeypatch.setattr(litellm_module, "LITELLM_AVAILABLE", True) - monkeypatch.setattr( - litellm_module, - "litellm_get_model_info", - lambda model: (_ for _ in ()).throw(RuntimeError("boom")), - ) - - provider = create_litellm_provider() - - assert isinstance(provider, LiteLLMProvider) - assert provider.get_context_limit("gpt-4o") == 128000 - assert provider.get_output_buffer("gpt-4o", default=3333) == 3333 - - @pytest.mark.skipif( - not is_litellm_available(), - reason="LiteLLM not installed", - ) - def test_create_litellm_provider(self): - """Test creating LiteLLM provider.""" - from headroom.providers import create_litellm_provider - - provider = create_litellm_provider() - assert provider.name == "litellm" - - @pytest.mark.skipif( - not is_litellm_available(), - reason="LiteLLM not installed", - ) - def test_litellm_supports_any_model(self): - """Test LiteLLM supports any model.""" - from headroom.providers import create_litellm_provider - - provider = create_litellm_provider() - assert provider.supports_model("gpt-4o") is True - assert provider.supports_model("claude-3-sonnet") is True - assert provider.supports_model("any-model") is True - - @pytest.mark.skipif( - not is_litellm_available(), - reason="LiteLLM not installed", - ) - def test_litellm_list_providers(self): - """Test listing LiteLLM providers.""" - from headroom.providers import LiteLLMProvider - - providers = LiteLLMProvider.list_supported_providers() - assert "openai" in providers - assert "anthropic" in providers - assert "ollama" in providers +"""Tests for universal provider support. + +Tests OpenAICompatibleProvider, GoogleProvider, and LiteLLMProvider. +""" + +from __future__ import annotations + +import pytest + +from headroom.providers import ( + GoogleProvider, + LiteLLMProvider, + ModelCapabilities, + OpenAICompatibleProvider, + create_anyscale_provider, + create_fireworks_provider, + create_groq_provider, + create_litellm_provider, + create_lmstudio_provider, + create_ollama_provider, + create_together_provider, + create_vllm_provider, + is_litellm_available, +) + + +def _transformers_available() -> bool: + """Check if transformers is available.""" + try: + import transformers # noqa: F401 + + return True + except ImportError: + return False + + +class TestOpenAICompatibleProvider: + """Tests for OpenAICompatibleProvider.""" + + def test_init_default(self): + """Test initialization with defaults.""" + provider = OpenAICompatibleProvider() + assert provider.name == "openai_compatible" + assert provider.base_url is None + + def test_init_with_config(self): + """Test initialization with configuration.""" + provider = OpenAICompatibleProvider( + name="custom", + base_url="http://localhost:8080/v1", + api_key="test-key", + ) + assert provider.name == "custom" + assert provider.base_url == "http://localhost:8080/v1" + assert provider.api_key == "test-key" + + def test_supports_any_model(self): + """Test that provider supports any model.""" + provider = OpenAICompatibleProvider() + assert provider.supports_model("any-model") is True + assert provider.supports_model("llama-3") is True + assert provider.supports_model("custom-finetuned") is True + + @pytest.mark.skipif( + not _transformers_available(), + reason="transformers not installed - needed for HuggingFace tokenizer", + ) + def test_get_token_counter(self): + """Test getting token counter.""" + provider = OpenAICompatibleProvider() + counter = provider.get_token_counter("llama-3-8b") + assert counter is not None + # Should be able to count tokens + count = counter.count_text("Hello, world!") + assert count > 0 + + def test_get_context_limit_known_model(self): + """Test context limit for known models.""" + provider = OpenAICompatibleProvider() + # Llama 3.1 has 128K context + limit = provider.get_context_limit("llama-3.1-8b") + assert limit == 128000 + + def test_get_context_limit_unknown_model(self): + """Test context limit for unknown models (defaults to 128K).""" + provider = OpenAICompatibleProvider() + limit = provider.get_context_limit("unknown-model") + assert limit == 128000 + + def test_register_model(self): + """Test registering a custom model.""" + provider = OpenAICompatibleProvider() + provider.register_model( + "my-model", + context_window=64000, + max_output_tokens=8192, + input_cost_per_1m=1.0, + output_cost_per_1m=2.0, + ) + assert provider.get_context_limit("my-model") == 64000 + + def test_estimate_cost_registered_model(self): + """Test cost estimation for registered model.""" + provider = OpenAICompatibleProvider() + provider.register_model( + "priced-model", + input_cost_per_1m=1.0, + output_cost_per_1m=2.0, + ) + cost = provider.estimate_cost( + input_tokens=1000000, + output_tokens=500000, + model="priced-model", + ) + assert cost == 2.0 # 1.0 + 1.0 + + def test_estimate_cost_unknown_model(self): + """Test cost estimation returns None for unknown model.""" + provider = OpenAICompatibleProvider() + cost = provider.estimate_cost( + input_tokens=1000, + output_tokens=500, + model="unknown-model", + ) + assert cost is None + + def test_register_model_accepts_capabilities_object(self): + provider = OpenAICompatibleProvider() + caps = ModelCapabilities(model="caps-model", context_window=16000, tokenizer_backend="test") + + provider.register_model("caps-model", capabilities=caps) + + assert provider.get_context_limit("caps-model") == 16000 + + def test_get_token_counter_uses_registered_tokenizer_backend(self, monkeypatch): + recorded: list[tuple[str, str | None]] = [] + + class DummyTokenizer: + def count_text(self, text: str) -> int: + return len(text.split()) + + monkeypatch.setattr( + "headroom.providers.openai_compatible.get_tokenizer", + lambda model, backend=None: recorded.append((model, backend)) or DummyTokenizer(), + ) + provider = OpenAICompatibleProvider( + models={ + "custom-model": ModelCapabilities( + model="custom-model", + tokenizer_backend="custom-backend", + ) + } + ) + + counter = provider.get_token_counter("custom-model") + + assert counter.count_text("one two three") == 3 + assert recorded == [("custom-model", "custom-backend")] + + def test_openai_compatible_token_counter_counts_message_parts(self, monkeypatch): + class DummyTokenizer: + def count_text(self, text: str) -> int: + return len(text) + + monkeypatch.setattr( + "headroom.providers.openai_compatible.get_tokenizer", + lambda model, backend=None: DummyTokenizer(), + ) + counter = OpenAICompatibleProvider().get_token_counter("demo-model") + + tokens = counter.count_message( + { + "role": "user", + "content": [{"type": "text", "text": "hi"}, "there"], + "name": "tester", + "tool_calls": [{"function": {"name": "lookup", "arguments": '{"x":1}'}}], + "tool_call_id": "call_123", + } + ) + total = counter.count_messages( + [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": ["world"]}, + ] + ) + + assert tokens == 55 + assert total == 34 + + def test_openai_compatible_token_counter_ignores_unhandled_content_shapes(self, monkeypatch): + class DummyTokenizer: + def count_text(self, text: str) -> int: + return len(text) + + monkeypatch.setattr( + "headroom.providers.openai_compatible.get_tokenizer", + lambda model, backend=None: DummyTokenizer(), + ) + counter = OpenAICompatibleProvider().get_token_counter("demo-model") + + assert counter.count_message({"role": "user", "content": {}}) == 8 + assert counter.count_message({"role": "user", "content": [{"type": "image"}, 123]}) == 8 + + def test_get_context_limit_prefix_output_buffer_and_partial_pricing(self): + provider = OpenAICompatibleProvider( + models={ + "buffered": ModelCapabilities( + model="buffered", + max_output_tokens=1200, + input_cost_per_1m=1.0, + ) + } + ) + + assert provider.get_context_limit("mistral-custom") == 32768 + assert provider.get_output_buffer("buffered", default=4000) == 1200 + assert provider.get_output_buffer("unknown", default=2222) == 2222 + assert provider.estimate_cost(1000, 1000, "buffered") is None + + +class TestModelCapabilities: + """Tests for ModelCapabilities dataclass.""" + + def test_default_values(self): + """Test default capability values.""" + caps = ModelCapabilities(model="test-model") + assert caps.context_window == 128000 + assert caps.max_output_tokens == 4096 + assert caps.supports_tools is True + assert caps.supports_vision is False + assert caps.supports_streaming is True + + def test_custom_values(self): + """Test custom capability values.""" + caps = ModelCapabilities( + model="custom-model", + context_window=32000, + max_output_tokens=16384, + supports_tools=False, + supports_vision=True, + input_cost_per_1m=0.5, + output_cost_per_1m=1.5, + ) + assert caps.context_window == 32000 + assert caps.max_output_tokens == 16384 + assert caps.supports_tools is False + assert caps.supports_vision is True + assert caps.input_cost_per_1m == 0.5 + assert caps.output_cost_per_1m == 1.5 + + +class TestGoogleProvider: + """Tests for GoogleProvider.""" + + @pytest.fixture + def provider(self): + """Create Google provider.""" + return GoogleProvider() + + def test_name(self, provider): + """Test provider name.""" + assert provider.name == "google" + + def test_supports_gemini_models(self, provider): + """Test support for Gemini models.""" + assert provider.supports_model("gemini-2.0-flash") is True + assert provider.supports_model("gemini-1.5-pro") is True + assert provider.supports_model("gemini-1.5-flash") is True + + def test_not_supports_other_models(self, provider): + """Test non-support for other models.""" + assert provider.supports_model("gpt-4o") is False + assert provider.supports_model("claude-3") is False + + def test_get_token_counter(self, provider): + """Test getting token counter.""" + counter = provider.get_token_counter("gemini-2.0-flash") + assert counter is not None + count = counter.count_text("Hello, world!") + assert count > 0 + + def test_get_context_limit_gemini_2(self, provider): + """Test context limit for Gemini 2.0.""" + limit = provider.get_context_limit("gemini-2.0-flash") + # LiteLLM returns 1048576 (2^20), fallback returns 1000000 + assert limit in (1000000, 1048576) # ~1M tokens + + def test_get_context_limit_gemini_1_5_pro(self, provider): + """Test context limit for Gemini 1.5 Pro (2M!).""" + limit = provider.get_context_limit("gemini-1.5-pro") + # LiteLLM returns 2097152 (2^21), fallback returns 2000000 + assert limit in (2000000, 2097152) # ~2M tokens! + + def test_estimate_cost(self, provider): + """Test cost estimation.""" + cost = provider.estimate_cost( + input_tokens=1000000, + output_tokens=500000, + model="gemini-2.0-flash", + ) + assert cost is not None + # 1M input * $0.10 + 0.5M output * $0.40 = $0.10 + $0.20 = $0.30 + assert abs(cost - 0.30) < 0.01 + + def test_openai_compatible_url(self): + """Test OpenAI-compatible URL.""" + url = GoogleProvider.get_openai_compatible_url("test-key") + assert "generativelanguage.googleapis.com" in url + + +class TestProviderFactoryFunctions: + """Tests for provider factory functions.""" + + def test_create_ollama_provider(self): + """Test creating Ollama provider.""" + provider = create_ollama_provider() + assert provider.name == "ollama" + assert provider.base_url == "http://localhost:11434/v1" + + def test_create_ollama_provider_custom_url(self): + """Test creating Ollama provider with custom URL.""" + provider = create_ollama_provider("http://192.168.1.100:11434/v1") + assert provider.base_url == "http://192.168.1.100:11434/v1" + + def test_create_together_provider(self): + """Test creating Together provider.""" + provider = create_together_provider() + assert provider.name == "together" + assert "together.xyz" in provider.base_url + + def test_create_groq_provider(self): + """Test creating Groq provider.""" + provider = create_groq_provider() + assert provider.name == "groq" + assert "groq.com" in provider.base_url + + def test_create_vllm_provider(self): + """Test creating vLLM provider.""" + provider = create_vllm_provider("http://localhost:8000/v1") + assert provider.name == "vllm" + assert provider.base_url == "http://localhost:8000/v1" + + def test_create_lmstudio_provider(self): + """Test creating LM Studio provider.""" + provider = create_lmstudio_provider() + assert provider.name == "lmstudio" + assert provider.base_url == "http://localhost:1234/v1" + + def test_create_fireworks_and_anyscale_providers(self): + fireworks = create_fireworks_provider(api_key="fireworks-key") + anyscale = create_anyscale_provider(api_key="anyscale-key") + + assert fireworks.name == "fireworks" + assert fireworks.base_url == "https://api.fireworks.ai/inference/v1" + assert fireworks.api_key == "fireworks-key" + assert anyscale.name == "anyscale" + assert anyscale.base_url == "https://api.endpoints.anyscale.com/v1" + assert anyscale.api_key == "anyscale-key" + + +class TestLiteLLMProvider: + """Tests for LiteLLM provider.""" + + def test_is_litellm_available(self): + """Test checking LiteLLM availability.""" + result = is_litellm_available() + assert isinstance(result, bool) + + def test_unavailable_litellm_paths(self, monkeypatch): + import headroom.providers.litellm as litellm_module + + monkeypatch.setattr(litellm_module, "LITELLM_AVAILABLE", False) + + assert litellm_module.is_litellm_available() is False + assert litellm_module.LiteLLMProvider.list_supported_providers() == [] + with pytest.raises(RuntimeError, match="LiteLLM is required"): + litellm_module.LiteLLMTokenCounter("gpt-4o") + with pytest.raises(RuntimeError, match="LiteLLM is required"): + litellm_module.LiteLLMProvider() + + def test_litellm_token_counter_fallback_paths(self, monkeypatch): + import headroom.providers.litellm as litellm_module + + class DummyFallback: + def count_text(self, text: str) -> int: + return len(text.split()) + + monkeypatch.setattr(litellm_module, "LITELLM_AVAILABLE", True) + monkeypatch.setattr( + litellm_module, + "litellm_token_counter", + lambda **kwargs: (_ for _ in ()).throw(RuntimeError("boom")), + ) + monkeypatch.setattr(litellm_module, "EstimatingTokenCounter", DummyFallback) + + counter = litellm_module.LiteLLMTokenCounter("gpt-4o") + + assert counter.count_text("") == 0 + assert counter.count_text("one two three") == 3 + assert counter.count_message({"content": "one two"}) == 6 + assert counter.count_messages([]) == 0 + assert counter.count_messages([{"content": "one two"}, {"content": "three"}]) == 14 + + def test_litellm_provider_info_and_cost_fallbacks(self, monkeypatch): + import headroom.providers.litellm as litellm_module + + monkeypatch.setattr(litellm_module, "LITELLM_AVAILABLE", True) + monkeypatch.setattr( + litellm_module, + "litellm_get_model_info", + lambda model: { + "ctx-model": {"max_input_tokens": 64000}, + "max-model": {"max_tokens": 32000}, + "none-model": {"max_input_tokens": None, "max_output_tokens": None}, + "output-model": {"max_output_tokens": 6000}, + }[model], + ) + monkeypatch.setattr( + litellm_module, + "litellm", + type( + "LiteLLM", + (), + { + "completion_cost": staticmethod( + lambda **kwargs: 1.23 + if kwargs["model"] == "priced-model" + else (_ for _ in ()).throw(RuntimeError("missing price")) + ) + }, + )(), + ) + + provider = litellm_module.LiteLLMProvider() + + assert provider.get_context_limit("ctx-model") == 64000 + assert provider.get_context_limit("max-model") == 32000 + assert provider.get_context_limit("none-model") == 128000 + assert provider.get_output_buffer("output-model", default=4000) == 4000 + assert provider.get_output_buffer("none-model", default=2222) == 2222 + assert provider.estimate_cost(1000, 1000, "priced-model") == 1.23 + assert provider.estimate_cost(1000, 1000, "missing-price") is None + + def test_litellm_provider_handles_info_exceptions_and_factory(self, monkeypatch): + import headroom.providers.litellm as litellm_module + + monkeypatch.setattr(litellm_module, "LITELLM_AVAILABLE", True) + monkeypatch.setattr( + litellm_module, + "litellm_get_model_info", + lambda model: (_ for _ in ()).throw(RuntimeError("boom")), + ) + + provider = create_litellm_provider() + + assert isinstance(provider, LiteLLMProvider) + assert provider.get_context_limit("gpt-4o") == 128000 + assert provider.get_output_buffer("gpt-4o", default=3333) == 3333 + + @pytest.mark.skipif( + not is_litellm_available(), + reason="LiteLLM not installed", + ) + def test_create_litellm_provider(self): + """Test creating LiteLLM provider.""" + from headroom.providers import create_litellm_provider + + provider = create_litellm_provider() + assert provider.name == "litellm" + + @pytest.mark.skipif( + not is_litellm_available(), + reason="LiteLLM not installed", + ) + def test_litellm_supports_any_model(self): + """Test LiteLLM supports any model.""" + from headroom.providers import create_litellm_provider + + provider = create_litellm_provider() + assert provider.supports_model("gpt-4o") is True + assert provider.supports_model("claude-3-sonnet") is True + assert provider.supports_model("any-model") is True + + @pytest.mark.skipif( + not is_litellm_available(), + reason="LiteLLM not installed", + ) + def test_litellm_list_providers(self): + """Test listing LiteLLM providers.""" + from headroom.providers import LiteLLMProvider + + providers = LiteLLMProvider.list_supported_providers() + assert "openai" in providers + assert "anthropic" in providers + assert "ollama" in providers