mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
test: remove provider diff churn
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
JerrettDavis
parent
f7e3450381
commit
4576f9caba
4 changed files with 1386 additions and 1386 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue