test: model registry unit test suite (45 tests)
Covers model_registry.py without requiring a running service or LLM:
Empty/fresh state: no files, missing user dir
Save/load: round-trip, corrupt file fallback
Migration: v1 hosts/models, v1 no active, v0 flat, v0 empty url,
distill_backend_mid=local → distill role, saves file after migrate
Built-in resolution: claude_cli, gemini_api, gemini_cli, unknown → None
User model resolution: local_openai merges host, missing host → None
get_model_for_role: registry primary, built-in from registry, skips missing,
walks full backup chain, .env fallback, hardcoded fallback,
custom roles
get_best_local_model: prefers role chain, falls back to first local, None if no local
Host CRUD: create, update, unknown ID creates new, remove + cascades to models
Model CRUD: create, update, remove + clears role refs
set_role: assign model, assign built-in, clear with None, invalid slot,
unknown model ID, creates new role key
get_defined_roles: returns all settings roles, fills gaps with {}
Multi-user isolation: registries don't bleed across users
All tests use tmp_path + patch.object(config.settings, ...) — no real files touched.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Scott Idem
805
cortex/tests/test_model_registry.py
Normal file
805
cortex/tests/test_model_registry.py
Normal file
@@ -0,0 +1,805 @@
|
||||
"""
|
||||
Unit tests for model_registry.py — no HTTP, no LLM calls, no running service.
|
||||
|
||||
All file I/O is redirected to tmp_path via patch.object(config.settings, "home_dir", ...).
|
||||
|
||||
Coverage:
|
||||
- Empty registry (no files)
|
||||
- Save/load round-trip
|
||||
- Migration from local_llm.json (v0 flat and v1 hosts/models)
|
||||
- Host CRUD
|
||||
- Model CRUD (including role reference cleanup on remove)
|
||||
- Role assignment (set_role, validation)
|
||||
- Model resolution (_resolve_model: built-ins, local_openai, missing host/model)
|
||||
- get_model_for_role: registry chain → .env fallback → hardcoded fallback
|
||||
- get_best_local_model: role chain, first-local fallback, no-local case
|
||||
- Backup chain: skips missing models, returns next valid
|
||||
"""
|
||||
|
||||
import json
|
||||
import pytest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _home(tmp_path: Path) -> Path:
|
||||
"""Create a minimal home directory and return the root."""
|
||||
root = tmp_path / "home"
|
||||
root.mkdir()
|
||||
return root
|
||||
|
||||
|
||||
def _user_dir(home: Path, username: str = "scott") -> Path:
|
||||
d = home / username
|
||||
d.mkdir(exist_ok=True)
|
||||
return d
|
||||
|
||||
|
||||
def _write_registry(home: Path, data: dict, username: str = "scott") -> Path:
|
||||
_user_dir(home, username)
|
||||
path = home / username / "model_registry.json"
|
||||
path.write_text(json.dumps(data))
|
||||
return path
|
||||
|
||||
|
||||
def _write_local_llm(home: Path, data: dict, username: str = "scott") -> Path:
|
||||
_user_dir(home, username)
|
||||
path = home / username / "local_llm.json"
|
||||
path.write_text(json.dumps(data))
|
||||
return path
|
||||
|
||||
|
||||
def _read_registry(home: Path, username: str = "scott") -> dict:
|
||||
path = home / username / "model_registry.json"
|
||||
return json.loads(path.read_text())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Empty / fresh state
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_empty_registry_no_files(tmp_path):
|
||||
"""With no files, _load returns an empty structure."""
|
||||
home = _home(tmp_path)
|
||||
_user_dir(home)
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
data = reg._load("scott")
|
||||
assert data["version"] == 1
|
||||
assert data["hosts"] == []
|
||||
assert data["models"] == []
|
||||
assert data["roles"] == {}
|
||||
|
||||
|
||||
def test_empty_registry_missing_user_dir(tmp_path):
|
||||
"""Even with no user dir, _load returns an empty structure gracefully."""
|
||||
home = _home(tmp_path)
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
data = reg._load("nobody")
|
||||
assert data["hosts"] == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Save / load round-trip
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_save_and_load(tmp_path):
|
||||
home = _home(tmp_path)
|
||||
_user_dir(home)
|
||||
import config
|
||||
import model_registry as reg
|
||||
|
||||
registry = {
|
||||
"version": 1,
|
||||
"hosts": [{"id": "h1", "label": "ML Box", "api_url": "http://10.0.0.1:3000", "api_key": "sk-test"}],
|
||||
"models": [{"id": "m1", "type": "local_openai", "label": "Gemma Small",
|
||||
"model_name": "gemma4:e4b", "host_id": "h1", "context_k": 72, "tags": ["fast"]}],
|
||||
"roles": {"chat": {"primary": "m1"}},
|
||||
}
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
reg._save("scott", registry)
|
||||
loaded = reg._load("scott")
|
||||
|
||||
assert loaded["hosts"][0]["label"] == "ML Box"
|
||||
assert loaded["models"][0]["model_name"] == "gemma4:e4b"
|
||||
assert loaded["roles"]["chat"]["primary"] == "m1"
|
||||
|
||||
|
||||
def test_corrupt_registry_falls_back_to_empty(tmp_path):
|
||||
home = _home(tmp_path)
|
||||
path = _user_dir(home) / "model_registry.json"
|
||||
path.write_text("{bad json{{")
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
data = reg._load("scott")
|
||||
assert data["hosts"] == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Migration from local_llm.json
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_migrate_v1_hosts_models(tmp_path):
|
||||
"""v1 local_llm.json (hosts/models/active_model_id) migrates correctly."""
|
||||
home = _home(tmp_path)
|
||||
_write_local_llm(home, {
|
||||
"hosts": [{"id": "h1", "label": "Home", "api_url": "http://10.0.0.1:3000", "api_key": "sk-1"}],
|
||||
"models": [
|
||||
{"id": "m1", "host_id": "h1", "label": "Gemma Small", "model_name": "gemma4:e4b"},
|
||||
{"id": "m2", "host_id": "h1", "label": "Gemma Med", "model_name": "gemma4:26b"},
|
||||
],
|
||||
"active_model_id": "m1",
|
||||
})
|
||||
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
data = reg._load("scott")
|
||||
|
||||
assert len(data["hosts"]) == 1
|
||||
assert data["hosts"][0]["api_url"] == "http://10.0.0.1:3000"
|
||||
assert len(data["models"]) == 2
|
||||
assert all(m["type"] == "local_openai" for m in data["models"])
|
||||
# active_model_id → roles.chat.primary
|
||||
assert data["roles"].get("chat", {}).get("primary") == "m1"
|
||||
|
||||
|
||||
def test_migrate_v1_no_active_model(tmp_path):
|
||||
"""Migration with active_model_id=null: chat role stays unset."""
|
||||
home = _home(tmp_path)
|
||||
_write_local_llm(home, {
|
||||
"hosts": [{"id": "h1", "label": "Box", "api_url": "http://10.0.0.1:3000", "api_key": ""}],
|
||||
"models": [{"id": "m1", "host_id": "h1", "label": "Model", "model_name": "llama3"}],
|
||||
"active_model_id": None,
|
||||
})
|
||||
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
data = reg._load("scott")
|
||||
|
||||
assert "chat" not in data["roles"] or data["roles"]["chat"].get("primary") is None
|
||||
|
||||
|
||||
def test_migrate_v0_flat_format(tmp_path):
|
||||
"""v0 flat local_llm.json is wrapped into hosts/models structure."""
|
||||
home = _home(tmp_path)
|
||||
_write_local_llm(home, {
|
||||
"api_url": "http://10.0.0.2:3000",
|
||||
"api_key": "sk-flat",
|
||||
"model": "qwen3:8b",
|
||||
})
|
||||
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
data = reg._load("scott")
|
||||
|
||||
assert len(data["hosts"]) == 1
|
||||
assert data["hosts"][0]["api_url"] == "http://10.0.0.2:3000"
|
||||
assert len(data["models"]) == 1
|
||||
assert data["models"][0]["model_name"] == "qwen3:8b"
|
||||
|
||||
|
||||
def test_migrate_v0_empty_url_returns_empty(tmp_path):
|
||||
"""v0 with no api_url and no .env fallback → nothing to migrate, empty registry."""
|
||||
home = _home(tmp_path)
|
||||
_write_local_llm(home, {"api_url": "", "api_key": "", "model": ""})
|
||||
|
||||
import config
|
||||
import model_registry as reg
|
||||
with (
|
||||
patch.object(config.settings, "home_dir", home),
|
||||
patch.object(config.settings, "local_api_url", ""), # ensure no .env fallback
|
||||
patch.object(config.settings, "local_model", ""),
|
||||
):
|
||||
data = reg._load("scott")
|
||||
|
||||
assert data["hosts"] == []
|
||||
assert data["models"] == []
|
||||
|
||||
|
||||
def test_migrate_v1_distill_local_sets_role(tmp_path):
|
||||
"""When DISTILL_BACKEND_MID=local and active model exists, distill role is set."""
|
||||
home = _home(tmp_path)
|
||||
_write_local_llm(home, {
|
||||
"hosts": [{"id": "h1", "label": "Box", "api_url": "http://10.0.0.1:3000", "api_key": ""}],
|
||||
"models": [{"id": "m1", "host_id": "h1", "label": "G", "model_name": "gemma4:e4b"}],
|
||||
"active_model_id": "m1",
|
||||
})
|
||||
|
||||
import config
|
||||
import model_registry as reg
|
||||
with (
|
||||
patch.object(config.settings, "home_dir", home),
|
||||
patch.object(config.settings, "distill_backend_mid", "local"),
|
||||
):
|
||||
data = reg._load("scott")
|
||||
|
||||
assert data["roles"].get("distill", {}).get("primary") == "m1"
|
||||
|
||||
|
||||
def test_migration_saves_registry_file(tmp_path):
|
||||
"""After migration, model_registry.json is written so next load skips migration."""
|
||||
home = _home(tmp_path)
|
||||
_write_local_llm(home, {
|
||||
"hosts": [{"id": "h1", "label": "Box", "api_url": "http://10.0.0.1:3000", "api_key": ""}],
|
||||
"models": [],
|
||||
"active_model_id": None,
|
||||
})
|
||||
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
reg._load("scott") # triggers migration + save
|
||||
# Second load should read model_registry.json, not re-run migration
|
||||
data2 = reg._load("scott")
|
||||
|
||||
assert (home / "scott" / "model_registry.json").exists()
|
||||
assert data2["version"] == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Built-in model resolution
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_builtin_claude_cli(tmp_path):
|
||||
home = _home(tmp_path)
|
||||
_user_dir(home)
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
result = reg._resolve_model(reg._empty(), "claude_cli")
|
||||
assert result is not None
|
||||
assert result["type"] == "claude_cli"
|
||||
assert result["id"] == "claude_cli"
|
||||
|
||||
|
||||
def test_builtin_gemini_api(tmp_path):
|
||||
home = _home(tmp_path)
|
||||
_user_dir(home)
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
result = reg._resolve_model(reg._empty(), "gemini_api")
|
||||
assert result["type"] == "gemini_api"
|
||||
|
||||
|
||||
def test_builtin_gemini_cli(tmp_path):
|
||||
home = _home(tmp_path)
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
result = reg._resolve_model(reg._empty(), "gemini_cli")
|
||||
assert result["type"] == "gemini_cli"
|
||||
|
||||
|
||||
def test_builtin_unknown_returns_none(tmp_path):
|
||||
home = _home(tmp_path)
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
result = reg._resolve_model(reg._empty(), "does_not_exist")
|
||||
assert result is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# User model resolution
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_resolve_local_openai_merges_host(tmp_path):
|
||||
"""local_openai model gets api_url and api_key merged from its host."""
|
||||
home = _home(tmp_path)
|
||||
registry = {
|
||||
"version": 1,
|
||||
"hosts": [{"id": "h1", "label": "Box", "api_url": "http://10.0.0.1:3000", "api_key": "sk-test"}],
|
||||
"models": [{"id": "m1", "type": "local_openai", "label": "G", "model_name": "gemma4:e4b",
|
||||
"host_id": "h1", "context_k": 72, "tags": []}],
|
||||
"roles": {},
|
||||
}
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
result = reg._resolve_model(registry, "m1")
|
||||
assert result["api_url"] == "http://10.0.0.1:3000"
|
||||
assert result["api_key"] == "sk-test"
|
||||
assert result["model_name"] == "gemma4:e4b"
|
||||
|
||||
|
||||
def test_resolve_local_openai_missing_host_returns_none(tmp_path):
|
||||
"""A model pointing to a non-existent host_id returns None."""
|
||||
home = _home(tmp_path)
|
||||
registry = {
|
||||
"version": 1, "hosts": [], "roles": {},
|
||||
"models": [{"id": "m1", "type": "local_openai", "host_id": "missing",
|
||||
"label": "X", "model_name": "x", "context_k": 0, "tags": []}],
|
||||
}
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
result = reg._resolve_model(registry, "m1")
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_resolve_unknown_model_id_returns_none(tmp_path):
|
||||
home = _home(tmp_path)
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
result = reg._resolve_model(reg._empty(), "no_such_model")
|
||||
assert result is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_model_for_role
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_get_model_for_role_uses_registry(tmp_path):
|
||||
"""Registry primary assignment is returned first."""
|
||||
home = _home(tmp_path)
|
||||
_write_registry(home, {
|
||||
"version": 1,
|
||||
"hosts": [{"id": "h1", "label": "Box", "api_url": "http://10.0.0.1:3000", "api_key": ""}],
|
||||
"models": [{"id": "m1", "type": "local_openai", "label": "G",
|
||||
"model_name": "gemma4:e4b", "host_id": "h1", "context_k": 72, "tags": []}],
|
||||
"roles": {"chat": {"primary": "m1"}},
|
||||
})
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
result = reg.get_model_for_role("scott", "chat")
|
||||
assert result["model_name"] == "gemma4:e4b"
|
||||
assert result["api_url"] == "http://10.0.0.1:3000"
|
||||
|
||||
|
||||
def test_get_model_for_role_uses_builtin_from_registry(tmp_path):
|
||||
"""Registry can assign built-in IDs (claude_cli, gemini_api, etc.)."""
|
||||
home = _home(tmp_path)
|
||||
_write_registry(home, {
|
||||
"version": 1, "hosts": [], "models": [],
|
||||
"roles": {"chat": {"primary": "claude_cli"}},
|
||||
})
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
result = reg.get_model_for_role("scott", "chat")
|
||||
assert result["type"] == "claude_cli"
|
||||
|
||||
|
||||
def test_get_model_for_role_skips_missing_primary(tmp_path):
|
||||
"""If primary model_id is not found, falls through to backup_1."""
|
||||
home = _home(tmp_path)
|
||||
_write_registry(home, {
|
||||
"version": 1,
|
||||
"hosts": [{"id": "h1", "label": "Box", "api_url": "http://10.0.0.1:3000", "api_key": ""}],
|
||||
"models": [{"id": "m2", "type": "local_openai", "label": "Backup",
|
||||
"model_name": "llama3:8b", "host_id": "h1", "context_k": 8, "tags": []}],
|
||||
"roles": {"chat": {"primary": "gone", "backup_1": "m2"}},
|
||||
})
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
result = reg.get_model_for_role("scott", "chat")
|
||||
assert result["model_name"] == "llama3:8b"
|
||||
|
||||
|
||||
def test_get_model_for_role_env_fallback(tmp_path):
|
||||
"""No registry entry for role → falls back to .env setting."""
|
||||
home = _home(tmp_path)
|
||||
_user_dir(home)
|
||||
import config
|
||||
import model_registry as reg
|
||||
with (
|
||||
patch.object(config.settings, "home_dir", home),
|
||||
patch.object(config.settings, "role_chat", "gemini_cli"),
|
||||
):
|
||||
result = reg.get_model_for_role("scott", "chat")
|
||||
assert result["type"] == "gemini_cli"
|
||||
|
||||
|
||||
def test_get_model_for_role_hardcoded_fallback(tmp_path):
|
||||
"""No registry + no .env for role → hardcoded last resort."""
|
||||
home = _home(tmp_path)
|
||||
_user_dir(home)
|
||||
import config
|
||||
import model_registry as reg
|
||||
# Clear the .env default for 'chat' to simulate unset
|
||||
with (
|
||||
patch.object(config.settings, "home_dir", home),
|
||||
patch.object(config.settings, "role_chat", ""),
|
||||
):
|
||||
result = reg.get_model_for_role("scott", "chat")
|
||||
# claude_cli is the hardcoded last resort for 'chat'
|
||||
assert result["type"] == "claude_cli"
|
||||
|
||||
|
||||
def test_get_model_for_role_custom_role(tmp_path):
|
||||
"""Custom roles not in DEFINED_ROLES can still be assigned and resolved."""
|
||||
home = _home(tmp_path)
|
||||
_write_registry(home, {
|
||||
"version": 1, "hosts": [], "models": [],
|
||||
"roles": {"therapy": {"primary": "gemini_api"}},
|
||||
})
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
result = reg.get_model_for_role("scott", "therapy")
|
||||
assert result["type"] == "gemini_api"
|
||||
|
||||
|
||||
def test_get_model_for_role_full_backup_chain(tmp_path):
|
||||
"""Walks the entire priority chain before falling back."""
|
||||
home = _home(tmp_path)
|
||||
_write_registry(home, {
|
||||
"version": 1,
|
||||
"hosts": [{"id": "h1", "label": "Box", "api_url": "http://10.0.0.1:3000", "api_key": ""}],
|
||||
"models": [{"id": "m4", "type": "local_openai", "label": "Last",
|
||||
"model_name": "tiny:1b", "host_id": "h1", "context_k": 4, "tags": []}],
|
||||
"roles": {"chat": {
|
||||
"primary": "gone1",
|
||||
"backup_1": "gone2",
|
||||
"backup_2": "gone3",
|
||||
"backup_3": "gone4",
|
||||
"backup_4": "m4",
|
||||
}},
|
||||
})
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
result = reg.get_model_for_role("scott", "chat")
|
||||
assert result["model_name"] == "tiny:1b"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_best_local_model
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_get_best_local_prefers_role_chain(tmp_path):
|
||||
"""Returns the first local_openai model in the chat role chain."""
|
||||
home = _home(tmp_path)
|
||||
_write_registry(home, {
|
||||
"version": 1,
|
||||
"hosts": [{"id": "h1", "label": "Box", "api_url": "http://10.0.0.1:3000", "api_key": ""}],
|
||||
"models": [
|
||||
{"id": "m1", "type": "local_openai", "label": "Preferred",
|
||||
"model_name": "gemma4:e4b", "host_id": "h1", "context_k": 72, "tags": []},
|
||||
],
|
||||
"roles": {"chat": {"primary": "claude_cli", "backup_1": "m1"}},
|
||||
})
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
# primary is claude_cli (not local), backup_1 is m1 (local)
|
||||
result = reg.get_best_local_model("scott", "chat")
|
||||
assert result["model_name"] == "gemma4:e4b"
|
||||
|
||||
|
||||
def test_get_best_local_falls_back_to_first_model(tmp_path):
|
||||
"""No local model in role chain → returns first configured local model."""
|
||||
home = _home(tmp_path)
|
||||
_write_registry(home, {
|
||||
"version": 1,
|
||||
"hosts": [{"id": "h1", "label": "Box", "api_url": "http://10.0.0.1:3000", "api_key": ""}],
|
||||
"models": [
|
||||
{"id": "m1", "type": "local_openai", "label": "G",
|
||||
"model_name": "gemma4:e4b", "host_id": "h1", "context_k": 72, "tags": []},
|
||||
],
|
||||
"roles": {}, # no chat role assigned
|
||||
})
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
result = reg.get_best_local_model("scott", "chat")
|
||||
assert result["model_name"] == "gemma4:e4b"
|
||||
|
||||
|
||||
def test_get_best_local_returns_none_when_no_local_models(tmp_path):
|
||||
"""No local_openai models configured → returns None."""
|
||||
home = _home(tmp_path)
|
||||
_write_registry(home, {
|
||||
"version": 1, "hosts": [], "models": [],
|
||||
"roles": {"chat": {"primary": "claude_cli"}},
|
||||
})
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
result = reg.get_best_local_model("scott", "chat")
|
||||
assert result is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Host CRUD
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_save_host_creates_new(tmp_path):
|
||||
home = _home(tmp_path)
|
||||
_user_dir(home)
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
host_id = reg.save_host("scott", None, "ML Box", "http://10.0.0.1:3000", "sk-abc")
|
||||
data = reg._load("scott")
|
||||
assert len(data["hosts"]) == 1
|
||||
assert data["hosts"][0]["id"] == host_id
|
||||
assert data["hosts"][0]["label"] == "ML Box"
|
||||
assert data["hosts"][0]["api_key"] == "sk-abc"
|
||||
|
||||
|
||||
def test_save_host_updates_existing(tmp_path):
|
||||
home = _home(tmp_path)
|
||||
_write_registry(home, {
|
||||
"version": 1,
|
||||
"hosts": [{"id": "h1", "label": "Old Label", "api_url": "http://10.0.0.1:3000", "api_key": "sk-old"}],
|
||||
"models": [], "roles": {},
|
||||
})
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
reg.save_host("scott", "h1", "New Label", "http://10.0.0.2:3000", "")
|
||||
data = reg._load("scott")
|
||||
assert len(data["hosts"]) == 1
|
||||
assert data["hosts"][0]["label"] == "New Label"
|
||||
assert data["hosts"][0]["api_url"] == "http://10.0.0.2:3000"
|
||||
# Empty api_key → existing key preserved
|
||||
assert data["hosts"][0]["api_key"] == "sk-old"
|
||||
|
||||
|
||||
def test_save_host_unknown_id_creates_new(tmp_path):
|
||||
"""Passing a host_id that doesn't exist creates a new host instead of crashing."""
|
||||
home = _home(tmp_path)
|
||||
_write_registry(home, {"version": 1, "hosts": [], "models": [], "roles": {}})
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
reg.save_host("scott", "ghost-id", "New", "http://10.0.0.3:3000", "")
|
||||
data = reg._load("scott")
|
||||
assert len(data["hosts"]) == 1
|
||||
|
||||
|
||||
def test_remove_host_also_removes_models(tmp_path):
|
||||
home = _home(tmp_path)
|
||||
_write_registry(home, {
|
||||
"version": 1,
|
||||
"hosts": [{"id": "h1", "label": "Box", "api_url": "http://10.0.0.1:3000", "api_key": ""}],
|
||||
"models": [{"id": "m1", "type": "local_openai", "host_id": "h1",
|
||||
"label": "G", "model_name": "gemma4:e4b", "context_k": 72, "tags": []}],
|
||||
"roles": {"chat": {"primary": "m1"}},
|
||||
})
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
found = reg.remove_host("scott", "h1")
|
||||
data = reg._load("scott")
|
||||
assert found is True
|
||||
assert data["hosts"] == []
|
||||
assert data["models"] == []
|
||||
|
||||
|
||||
def test_remove_host_not_found_returns_false(tmp_path):
|
||||
home = _home(tmp_path)
|
||||
_write_registry(home, {"version": 1, "hosts": [], "models": [], "roles": {}})
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
found = reg.remove_host("scott", "nope")
|
||||
assert found is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Model CRUD
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_save_model_creates(tmp_path):
|
||||
home = _home(tmp_path)
|
||||
_write_registry(home, {
|
||||
"version": 1,
|
||||
"hosts": [{"id": "h1", "label": "Box", "api_url": "http://10.0.0.1:3000", "api_key": ""}],
|
||||
"models": [], "roles": {},
|
||||
})
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
model_id = reg.save_model("scott", None, "h1", "Gemma Small", "gemma4:e4b", 72, ["fast", "distill"])
|
||||
data = reg._load("scott")
|
||||
assert len(data["models"]) == 1
|
||||
assert data["models"][0]["id"] == model_id
|
||||
assert data["models"][0]["context_k"] == 72
|
||||
assert data["models"][0]["tags"] == ["fast", "distill"]
|
||||
assert data["models"][0]["type"] == "local_openai"
|
||||
|
||||
|
||||
def test_save_model_updates_existing(tmp_path):
|
||||
home = _home(tmp_path)
|
||||
_write_registry(home, {
|
||||
"version": 1,
|
||||
"hosts": [{"id": "h1", "label": "Box", "api_url": "http://10.0.0.1:3000", "api_key": ""}],
|
||||
"models": [{"id": "m1", "type": "local_openai", "label": "Old",
|
||||
"model_name": "llama3", "host_id": "h1", "context_k": 8, "tags": []}],
|
||||
"roles": {},
|
||||
})
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
reg.save_model("scott", "m1", "h1", "New Label", "llama3:latest", 128, ["updated"])
|
||||
data = reg._load("scott")
|
||||
assert len(data["models"]) == 1
|
||||
assert data["models"][0]["label"] == "New Label"
|
||||
assert data["models"][0]["context_k"] == 128
|
||||
|
||||
|
||||
def test_remove_model_clears_role_refs(tmp_path):
|
||||
"""Removing a model clears it from any role assignments."""
|
||||
home = _home(tmp_path)
|
||||
_write_registry(home, {
|
||||
"version": 1,
|
||||
"hosts": [{"id": "h1", "label": "Box", "api_url": "http://10.0.0.1:3000", "api_key": ""}],
|
||||
"models": [{"id": "m1", "type": "local_openai", "label": "G",
|
||||
"model_name": "gemma4:e4b", "host_id": "h1", "context_k": 72, "tags": []}],
|
||||
"roles": {
|
||||
"chat": {"primary": "m1", "backup_1": "m1"},
|
||||
"distill": {"primary": "claude_cli", "backup_1": "m1"},
|
||||
},
|
||||
})
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
found = reg.remove_model("scott", "m1")
|
||||
data = reg._load("scott")
|
||||
assert found is True
|
||||
assert data["models"] == []
|
||||
assert data["roles"]["chat"].get("primary") is None
|
||||
assert data["roles"]["chat"].get("backup_1") is None
|
||||
assert data["roles"]["distill"].get("backup_1") is None
|
||||
# claude_cli assignment preserved
|
||||
assert data["roles"]["distill"]["primary"] == "claude_cli"
|
||||
|
||||
|
||||
def test_remove_model_not_found_returns_false(tmp_path):
|
||||
home = _home(tmp_path)
|
||||
_write_registry(home, {"version": 1, "hosts": [], "models": [], "roles": {}})
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
found = reg.remove_model("scott", "ghost")
|
||||
assert found is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# set_role
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_set_role_assigns_model(tmp_path):
|
||||
home = _home(tmp_path)
|
||||
_write_registry(home, {
|
||||
"version": 1,
|
||||
"hosts": [{"id": "h1", "label": "Box", "api_url": "http://10.0.0.1:3000", "api_key": ""}],
|
||||
"models": [{"id": "m1", "type": "local_openai", "label": "G",
|
||||
"model_name": "gemma4:e4b", "host_id": "h1", "context_k": 72, "tags": []}],
|
||||
"roles": {},
|
||||
})
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
ok = reg.set_role("scott", "chat", "primary", "m1")
|
||||
data = reg._load("scott")
|
||||
assert ok is True
|
||||
assert data["roles"]["chat"]["primary"] == "m1"
|
||||
|
||||
|
||||
def test_set_role_assigns_builtin(tmp_path):
|
||||
home = _home(tmp_path)
|
||||
_write_registry(home, {"version": 1, "hosts": [], "models": [], "roles": {}})
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
ok = reg.set_role("scott", "orchestrator", "primary", "gemini_api")
|
||||
data = reg._load("scott")
|
||||
assert ok is True
|
||||
assert data["roles"]["orchestrator"]["primary"] == "gemini_api"
|
||||
|
||||
|
||||
def test_set_role_clears_with_none(tmp_path):
|
||||
home = _home(tmp_path)
|
||||
_write_registry(home, {
|
||||
"version": 1, "hosts": [], "models": [],
|
||||
"roles": {"chat": {"primary": "claude_cli"}},
|
||||
})
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
ok = reg.set_role("scott", "chat", "primary", None)
|
||||
data = reg._load("scott")
|
||||
assert ok is True
|
||||
assert data["roles"]["chat"]["primary"] is None
|
||||
|
||||
|
||||
def test_set_role_invalid_slot_returns_false(tmp_path):
|
||||
home = _home(tmp_path)
|
||||
_write_registry(home, {"version": 1, "hosts": [], "models": [], "roles": {}})
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
ok = reg.set_role("scott", "chat", "backup_99", "claude_cli")
|
||||
assert ok is False
|
||||
|
||||
|
||||
def test_set_role_unknown_model_id_returns_false(tmp_path):
|
||||
home = _home(tmp_path)
|
||||
_write_registry(home, {"version": 1, "hosts": [], "models": [], "roles": {}})
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
ok = reg.set_role("scott", "chat", "primary", "nonexistent_model")
|
||||
assert ok is False
|
||||
|
||||
|
||||
def test_set_role_creates_role_key_if_missing(tmp_path):
|
||||
"""set_role on a role that isn't in roles{} yet creates it."""
|
||||
home = _home(tmp_path)
|
||||
_write_registry(home, {"version": 1, "hosts": [], "models": [], "roles": {}})
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
reg.set_role("scott", "medical", "primary", "claude_cli")
|
||||
data = reg._load("scott")
|
||||
assert data["roles"]["medical"]["primary"] == "claude_cli"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_defined_roles
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_get_defined_roles_returns_registry_roles(tmp_path):
|
||||
home = _home(tmp_path)
|
||||
_write_registry(home, {
|
||||
"version": 1, "hosts": [], "models": [],
|
||||
"roles": {"chat": {"primary": "claude_cli"}, "distill": {}},
|
||||
})
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
roles = reg.get_defined_roles("scott")
|
||||
# Should include all settings.defined_roles, filling gaps with {}
|
||||
for role in config.settings.get_defined_roles():
|
||||
assert role in roles
|
||||
|
||||
|
||||
def test_get_defined_roles_fills_gaps(tmp_path):
|
||||
"""Roles in settings.defined_roles that aren't in registry get empty dicts."""
|
||||
home = _home(tmp_path)
|
||||
_write_registry(home, {"version": 1, "hosts": [], "models": [], "roles": {}})
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
roles = reg.get_defined_roles("scott")
|
||||
assert "chat" in roles
|
||||
assert roles["chat"] == {}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Multi-user isolation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_registries_are_isolated_per_user(tmp_path):
|
||||
"""Each user has their own registry file — changes don't bleed across users."""
|
||||
home = _home(tmp_path)
|
||||
(home / "scott").mkdir()
|
||||
(home / "holly").mkdir()
|
||||
|
||||
import config
|
||||
import model_registry as reg
|
||||
with patch.object(config.settings, "home_dir", home):
|
||||
reg.save_host("scott", None, "Scott Host", "http://10.0.0.1:3000", "")
|
||||
scott_data = reg._load("scott")
|
||||
holly_data = reg._load("holly")
|
||||
|
||||
assert len(scott_data["hosts"]) == 1
|
||||
assert holly_data["hosts"] == []
|
||||
Reference in New Issue
Block a user