411 lines
16 KiB
Python
411 lines
16 KiB
Python
"""Unit tests for LlmConfigService (U5 — LLM provider/fallback CRUD).
|
|
|
|
Covers:
|
|
- Happy path: list → get → create → update → delete providers
|
|
- API key management: set_api_key writes to .env, masks in responses
|
|
- Fallback chain: get → set → delete
|
|
- Edge case: create duplicate provider → ValueError
|
|
- Edge case: update non-existent provider → ValueError
|
|
- Edge case: delete provider used in fallback → ValueError
|
|
- Edge case: delete non-existent provider → ValueError
|
|
- File locking: concurrent writes don't corrupt (sequential simulation)
|
|
- Singleton helpers (get/set_llm_config_service)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pytest
|
|
import yaml
|
|
|
|
from agentkit.server.admin.llm_config_service import (
|
|
LlmConfigService,
|
|
get_llm_config_service,
|
|
set_llm_config_service,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fixtures
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _sample_config() -> dict[str, Any]:
|
|
"""A minimal agentkit.yaml-style config for testing."""
|
|
return {
|
|
"server": {"host": "0.0.0.0", "port": 8001},
|
|
"llm": {
|
|
"providers": {
|
|
"openai": {
|
|
"type": "openai",
|
|
"api_key": "sk-test-12345678",
|
|
"base_url": "https://api.openai.com/v1",
|
|
"models": {"gpt-4o": {}},
|
|
"max_tokens": 4096,
|
|
"timeout": 120.0,
|
|
},
|
|
},
|
|
"model_aliases": {"gpt4": "openai/gpt-4o"},
|
|
"fallbacks": {},
|
|
},
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def config_path(tmp_path: Path) -> Path:
|
|
"""Create a temporary agentkit.yaml config file."""
|
|
path = tmp_path / "agentkit.yaml"
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
yaml.dump(_sample_config(), f, default_flow_style=False, allow_unicode=True)
|
|
return path
|
|
|
|
|
|
@pytest.fixture
|
|
def service(config_path: Path) -> LlmConfigService:
|
|
return LlmConfigService(config_path)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _reset_singleton():
|
|
"""Reset the module singleton before and after each test."""
|
|
set_llm_config_service(None)
|
|
yield
|
|
set_llm_config_service(None)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Provider CRUD happy path
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestProviderCrudHappyPath:
|
|
def test_list_returns_existing_providers(self, service: LlmConfigService):
|
|
providers = service.list_providers()
|
|
assert len(providers) == 1
|
|
assert providers[0]["name"] == "openai"
|
|
assert providers[0]["type"] == "openai"
|
|
|
|
def test_list_masks_api_keys(self, service: LlmConfigService):
|
|
providers = service.list_providers()
|
|
# API key should be masked: ****xxxx
|
|
assert providers[0]["api_key"].startswith("****")
|
|
# Should show only last 4 chars
|
|
assert providers[0]["api_key"] == "****5678"
|
|
|
|
def test_get_returns_provider_by_name(self, service: LlmConfigService):
|
|
provider = service.get_provider("openai")
|
|
assert provider is not None
|
|
assert provider["name"] == "openai"
|
|
assert provider["type"] == "openai"
|
|
assert provider["base_url"] == "https://api.openai.com/v1"
|
|
|
|
def test_get_returns_none_for_unknown(self, service: LlmConfigService):
|
|
assert service.get_provider("nonexistent") is None
|
|
|
|
def test_create_adds_provider_to_yaml(self, service: LlmConfigService, config_path: Path):
|
|
created = service.create_provider(
|
|
name="anthropic",
|
|
provider_type="anthropic",
|
|
api_key="sk-ant-test-abcd1234",
|
|
base_url="https://api.anthropic.com",
|
|
models={"claude-sonnet-4-20250514": {}},
|
|
max_tokens=8192,
|
|
timeout=90.0,
|
|
)
|
|
assert created["name"] == "anthropic"
|
|
assert created["type"] == "anthropic"
|
|
# API key should be masked in the response.
|
|
assert created["api_key"].startswith("****")
|
|
# Verify the YAML file was updated.
|
|
with open(config_path, encoding="utf-8") as f:
|
|
saved = yaml.safe_load(f)
|
|
assert "anthropic" in saved["llm"]["providers"]
|
|
# The YAML should store a ${VAR} reference, not the plaintext key.
|
|
assert saved["llm"]["providers"]["anthropic"]["api_key"] == "${ANTHROPIC_API_KEY}"
|
|
|
|
def test_create_writes_api_key_to_env_file(self, service: LlmConfigService, config_path: Path):
|
|
service.create_provider(
|
|
name="anthropic",
|
|
provider_type="anthropic",
|
|
api_key="sk-ant-test-abcd1234",
|
|
)
|
|
env_path = config_path.parent / ".env"
|
|
assert env_path.exists()
|
|
with open(env_path, encoding="utf-8") as f:
|
|
content = f.read()
|
|
assert "ANTHROPIC_API_KEY=sk-ant-test-abcd1234" in content
|
|
|
|
def test_update_partial_only_changes_provided_fields(
|
|
self, service: LlmConfigService, config_path: Path
|
|
):
|
|
updated = service.update_provider("openai", max_tokens=2048)
|
|
assert updated["max_tokens"] == 2048
|
|
# Other fields should be preserved.
|
|
assert updated["base_url"] == "https://api.openai.com/v1"
|
|
assert updated["type"] == "openai"
|
|
|
|
def test_update_with_masked_api_key_preserves_existing(
|
|
self, service: LlmConfigService, config_path: Path
|
|
):
|
|
"""When user sends back a masked key (****xxxx), the real key should
|
|
be preserved."""
|
|
service.update_provider("openai", api_key="****5678")
|
|
with open(config_path, encoding="utf-8") as f:
|
|
saved = yaml.safe_load(f)
|
|
# The original plaintext key should be preserved.
|
|
assert saved["llm"]["providers"]["openai"]["api_key"] == "sk-test-12345678"
|
|
|
|
def test_update_with_new_api_key_writes_to_env(
|
|
self, service: LlmConfigService, config_path: Path
|
|
):
|
|
service.update_provider("openai", api_key="sk-new-key-9999")
|
|
env_path = config_path.parent / ".env"
|
|
with open(env_path, encoding="utf-8") as f:
|
|
content = f.read()
|
|
assert "OPENAI_API_KEY=sk-new-key-9999" in content
|
|
# YAML should now have a ${VAR} reference.
|
|
with open(config_path, encoding="utf-8") as f:
|
|
saved = yaml.safe_load(f)
|
|
assert saved["llm"]["providers"]["openai"]["api_key"] == "${OPENAI_API_KEY}"
|
|
|
|
def test_delete_removes_provider(self, service: LlmConfigService, config_path: Path):
|
|
deleted = service.delete_provider("openai")
|
|
assert deleted is True
|
|
# Verify the YAML file was updated.
|
|
with open(config_path, encoding="utf-8") as f:
|
|
saved = yaml.safe_load(f)
|
|
assert "openai" not in saved["llm"]["providers"]
|
|
# list_providers should return empty.
|
|
assert service.list_providers() == []
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Provider CRUD edge cases
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestProviderCrudEdgeCases:
|
|
def test_create_duplicate_raises_value_error(self, service: LlmConfigService):
|
|
with pytest.raises(ValueError, match="already exists"):
|
|
service.create_provider(
|
|
name="openai",
|
|
provider_type="openai",
|
|
api_key="sk-duplicate",
|
|
)
|
|
|
|
def test_update_nonexistent_raises_value_error(self, service: LlmConfigService):
|
|
with pytest.raises(ValueError, match="not found"):
|
|
service.update_provider("nonexistent", max_tokens=2048)
|
|
|
|
def test_delete_nonexistent_raises_value_error(self, service: LlmConfigService):
|
|
with pytest.raises(ValueError, match="not found"):
|
|
service.delete_provider("nonexistent")
|
|
|
|
def test_delete_provider_used_in_fallback_raises_value_error(self, service: LlmConfigService):
|
|
# Add a second provider and a fallback chain that references it.
|
|
service.create_provider(
|
|
name="anthropic",
|
|
provider_type="anthropic",
|
|
api_key="sk-ant-test",
|
|
)
|
|
service.set_fallback("gpt-4o", ["anthropic"])
|
|
with pytest.raises(ValueError, match="used in fallback"):
|
|
service.delete_provider("anthropic")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# API key management
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestApiKeyManagement:
|
|
def test_set_api_key_writes_to_env_file(self, service: LlmConfigService, config_path: Path):
|
|
service.set_api_key("openai", "sk-brand-new-key-1234")
|
|
env_path = config_path.parent / ".env"
|
|
assert env_path.exists()
|
|
with open(env_path, encoding="utf-8") as f:
|
|
content = f.read()
|
|
assert "OPENAI_API_KEY=sk-brand-new-key-1234" in content
|
|
|
|
def test_set_api_key_updates_yaml_reference(self, service: LlmConfigService, config_path: Path):
|
|
service.set_api_key("openai", "sk-brand-new-key-1234")
|
|
with open(config_path, encoding="utf-8") as f:
|
|
saved = yaml.safe_load(f)
|
|
assert saved["llm"]["providers"]["openai"]["api_key"] == "${OPENAI_API_KEY}"
|
|
|
|
def test_set_api_key_returns_masked_provider(self, service: LlmConfigService):
|
|
result = service.set_api_key("openai", "sk-brand-new-key-1234")
|
|
assert result["api_key"].startswith("****")
|
|
# Should not contain the plaintext key.
|
|
assert "sk-brand-new-key-1234" not in result["api_key"]
|
|
|
|
def test_set_api_key_for_new_provider_creates_stub(
|
|
self, service: LlmConfigService, config_path: Path
|
|
):
|
|
"""Setting an API key for a provider that doesn't yet exist in the
|
|
YAML should create a stub entry."""
|
|
result = service.set_api_key("gemini", "AIza-test-key-1234")
|
|
assert result["name"] == "gemini"
|
|
assert result["type"] == "openai" # default type for stub
|
|
assert result["api_key"].startswith("****")
|
|
with open(config_path, encoding="utf-8") as f:
|
|
saved = yaml.safe_load(f)
|
|
assert "gemini" in saved["llm"]["providers"]
|
|
assert saved["llm"]["providers"]["gemini"]["api_key"] == "${GEMINI_API_KEY}"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fallback chain management
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestFallbackManagement:
|
|
def test_get_fallbacks_returns_empty_when_none_configured(self, service: LlmConfigService):
|
|
fallbacks = service.get_fallbacks()
|
|
assert fallbacks == {}
|
|
|
|
def test_set_fallback_adds_chain(self, service: LlmConfigService, config_path: Path):
|
|
result = service.set_fallback("gpt-4o", ["openai", "anthropic"])
|
|
assert result["model"] == "gpt-4o"
|
|
assert result["chain"] == ["openai", "anthropic"]
|
|
# Verify it was written to YAML.
|
|
with open(config_path, encoding="utf-8") as f:
|
|
saved = yaml.safe_load(f)
|
|
assert saved["llm"]["fallbacks"]["gpt-4o"] == ["openai", "anthropic"]
|
|
|
|
def test_get_fallbacks_returns_set_chain(self, service: LlmConfigService):
|
|
service.set_fallback("gpt-4o", ["openai", "anthropic"])
|
|
fallbacks = service.get_fallbacks()
|
|
assert fallbacks == {"gpt-4o": ["openai", "anthropic"]}
|
|
|
|
def test_set_fallback_overwrites_existing(self, service: LlmConfigService):
|
|
service.set_fallback("gpt-4o", ["openai"])
|
|
service.set_fallback("gpt-4o", ["anthropic", "gemini"])
|
|
fallbacks = service.get_fallbacks()
|
|
assert fallbacks["gpt-4o"] == ["anthropic", "gemini"]
|
|
|
|
def test_delete_fallback_removes_chain(self, service: LlmConfigService):
|
|
service.set_fallback("gpt-4o", ["openai"])
|
|
deleted = service.delete_fallback("gpt-4o")
|
|
assert deleted is True
|
|
assert service.get_fallbacks() == {}
|
|
|
|
def test_delete_fallback_returns_false_for_unset(self, service: LlmConfigService):
|
|
deleted = service.delete_fallback("nonexistent")
|
|
assert deleted is False
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# File locking
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestFileLocking:
|
|
def test_concurrent_writes_do_not_corrupt(self, service: LlmConfigService, config_path: Path):
|
|
"""Simulate concurrent writes from multiple threads.
|
|
|
|
Each thread creates a different provider. The file lock should
|
|
serialize the writes so the final YAML is valid and contains
|
|
all providers.
|
|
"""
|
|
errors: list[Exception] = []
|
|
|
|
def _create_provider(name: str, api_key: str) -> None:
|
|
try:
|
|
# Each thread needs its own service instance to avoid
|
|
# sharing state, but they all point to the same file.
|
|
svc = LlmConfigService(config_path)
|
|
svc.create_provider(name=name, provider_type="openai", api_key=api_key)
|
|
except Exception as exc:
|
|
errors.append(exc)
|
|
|
|
threads = [
|
|
threading.Thread(
|
|
target=_create_provider,
|
|
args=(f"provider_{i}", f"sk-key-{i:04d}"),
|
|
)
|
|
for i in range(5)
|
|
]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
# No errors should have occurred.
|
|
assert errors == [], f"Concurrent writes failed: {errors}"
|
|
|
|
# The YAML file should be valid and contain all 5 new providers
|
|
# plus the original "openai" provider.
|
|
with open(config_path, encoding="utf-8") as f:
|
|
saved = yaml.safe_load(f)
|
|
providers = saved["llm"]["providers"]
|
|
for i in range(5):
|
|
assert f"provider_{i}" in providers
|
|
assert "openai" in providers
|
|
|
|
def test_lockfile_is_created(self, service: LlmConfigService, config_path: Path):
|
|
"""The lockfile should be created on the first write."""
|
|
service.create_provider(
|
|
name="anthropic",
|
|
provider_type="anthropic",
|
|
api_key="sk-ant-test",
|
|
)
|
|
lockfile = config_path.with_suffix(config_path.suffix + ".lock")
|
|
assert lockfile.exists()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Singleton helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestSingletonHelpers:
|
|
def test_get_llm_config_service_requires_path_on_first_call(self):
|
|
with pytest.raises(ValueError, match="config_path is required"):
|
|
get_llm_config_service()
|
|
|
|
def test_get_llm_config_service_returns_singleton(self, config_path: Path):
|
|
svc1 = get_llm_config_service(config_path)
|
|
svc2 = get_llm_config_service() # subsequent call ignores path
|
|
assert svc1 is svc2
|
|
|
|
def test_set_llm_config_service_overrides_singleton(self, config_path: Path):
|
|
custom = LlmConfigService(config_path)
|
|
set_llm_config_service(custom)
|
|
assert get_llm_config_service() is custom
|
|
|
|
def test_set_llm_config_service_none_resets_singleton(self, config_path: Path):
|
|
svc = get_llm_config_service(config_path)
|
|
set_llm_config_service(None)
|
|
# Next call should raise (no path provided after reset).
|
|
with pytest.raises(ValueError, match="config_path is required"):
|
|
get_llm_config_service()
|
|
# The previously-created instance should still be usable.
|
|
assert svc.list_providers() is not None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Env var resolution
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestEnvVarResolution:
|
|
def test_list_providers_masks_env_var_referenced_keys(
|
|
self, config_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
):
|
|
"""When the YAML stores ``${VAR}``, the masked response should
|
|
show the masked *resolved* env value, not the literal ``${VAR}``."""
|
|
# Rewrite the config to use a ${VAR} reference.
|
|
with open(config_path, "w", encoding="utf-8") as f:
|
|
cfg = _sample_config()
|
|
cfg["llm"]["providers"]["openai"]["api_key"] = "${TEST_OPENAI_KEY}"
|
|
yaml.dump(cfg, f, default_flow_style=False, allow_unicode=True)
|
|
# Set the env var so it can be resolved.
|
|
monkeypatch.setenv("TEST_OPENAI_KEY", "sk-resolved-1234")
|
|
svc = LlmConfigService(config_path)
|
|
providers = svc.list_providers()
|
|
assert providers[0]["api_key"] == "****1234"
|