m4
405 lines
16 KiB
Python
405 lines
16 KiB
Python
"""Tests for the SQLite-backed ModelRuntimeStore (design doc 4.2, 5.1, 8.2)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
import stat
|
|
import threading
|
|
|
|
import pytest
|
|
|
|
from EvoScientist.model_registry.errors import (
|
|
CREDENTIAL_NOT_CONFIGURED,
|
|
MODEL_DISABLED,
|
|
RUN_CREDENTIAL_REVISION_UNAVAILABLE,
|
|
ModelRegistryError,
|
|
)
|
|
from EvoScientist.model_registry.schemas import CredentialWrite, RegistryV4
|
|
from EvoScientist.model_registry.store import ModelRuntimeStore
|
|
|
|
|
|
def _model_runtime(**overrides):
|
|
payload = {
|
|
"limit_mode": "combined",
|
|
"context_window_tokens": 1048576,
|
|
"max_input_tokens": None,
|
|
"max_output_tokens": 32768,
|
|
"min_effective_input_tokens": 8192,
|
|
"fixed_system_reserve_tokens": 4096,
|
|
"fixed_tools_reserve_tokens": 8192,
|
|
"fixed_attachments_reserve_tokens": 4096,
|
|
"limits_status": "confirmed",
|
|
"limits_source": "provider",
|
|
"temperature": None,
|
|
"top_p": None,
|
|
"reasoning_effort": "auto",
|
|
"declared_capabilities": {
|
|
"tools": True,
|
|
"vision": False,
|
|
"structured_output": True,
|
|
},
|
|
}
|
|
payload.update(overrides)
|
|
return payload
|
|
|
|
|
|
def _registry_payload(**overrides):
|
|
payload = {
|
|
"version": 4,
|
|
"revision": 1,
|
|
"state": "bootstrap",
|
|
"defaults": {
|
|
"primary": {"provider_id": "zhipu-glm", "model_key": "glm-5.2"},
|
|
},
|
|
"providers": [
|
|
{
|
|
"id": "zhipu-glm",
|
|
"name": "Zhipu GLM",
|
|
"adapter": "openai-compatible",
|
|
"base_url": "https://open.bigmodel.cn/api/paas/v4",
|
|
"auth": {"mode": "api_key", "credential_id": "zhipu-primary"},
|
|
"enabled": True,
|
|
"runtime": {"timeout_seconds": 120, "max_retries": 2},
|
|
"models": [
|
|
{
|
|
"key": "glm-5.2",
|
|
"name": "GLM-5.2",
|
|
"upstream_model_id": "glm-5.2",
|
|
"enabled": True,
|
|
"runtime": _model_runtime(),
|
|
}
|
|
],
|
|
}
|
|
],
|
|
}
|
|
payload.update(overrides)
|
|
return payload
|
|
|
|
|
|
def _registry(**overrides) -> RegistryV4:
|
|
return RegistryV4.model_validate(_registry_payload(**overrides))
|
|
|
|
|
|
def _local_registry() -> RegistryV4:
|
|
"""A registry whose provider needs no credential (mode=none)."""
|
|
payload = _registry_payload()
|
|
provider = payload["providers"][0]
|
|
provider["id"] = "local-ollama"
|
|
provider["adapter"] = "ollama"
|
|
provider["base_url"] = "http://localhost:11434"
|
|
provider["auth"] = {"mode": "none", "credential_id": None}
|
|
payload["defaults"]["primary"] = {
|
|
"provider_id": "local-ollama",
|
|
"model_key": "glm-5.2",
|
|
}
|
|
return RegistryV4.model_validate(payload)
|
|
|
|
|
|
class TestBootstrap:
|
|
def test_empty_database_loads_bootstrap_registry(self, tmp_path):
|
|
store = ModelRuntimeStore(config_dir=tmp_path)
|
|
registry = store.load_registry()
|
|
assert registry.version == 4
|
|
assert registry.revision == 1
|
|
assert registry.state == "bootstrap"
|
|
assert registry.providers == []
|
|
assert registry.defaults.primary is None
|
|
|
|
def test_database_path_and_permissions(self, tmp_path):
|
|
store = ModelRuntimeStore(config_dir=tmp_path)
|
|
assert store.database_path == tmp_path / "model-runtime.sqlite3"
|
|
file_mode = stat.S_IMODE(store.database_path.stat().st_mode)
|
|
dir_mode = stat.S_IMODE(tmp_path.stat().st_mode)
|
|
assert file_mode == 0o600
|
|
assert dir_mode == 0o700
|
|
|
|
def test_default_config_dir(self, tmp_path, monkeypatch):
|
|
from EvoScientist.model_registry import store as store_module
|
|
|
|
fake_home = tmp_path / "home"
|
|
monkeypatch.setattr(
|
|
store_module, "DEFAULT_CONFIG_DIR", fake_home / ".config" / "evoscientist"
|
|
)
|
|
store = ModelRuntimeStore()
|
|
assert store.database_path.parent.name == "evoscientist"
|
|
assert store.database_path.parent.parent.name == ".config"
|
|
assert store.database_path.exists()
|
|
|
|
def test_check_shared_storage_succeeds_on_local_volume(self, tmp_path):
|
|
store = ModelRuntimeStore(config_dir=tmp_path)
|
|
store.check_shared_storage()
|
|
|
|
|
|
class TestRegistrySave:
|
|
def test_first_ready_save_transitions_to_active(self, tmp_path):
|
|
store = ModelRuntimeStore(config_dir=tmp_path)
|
|
saved = store.save_registry(
|
|
expected_revision=1,
|
|
registry=_local_registry(),
|
|
)
|
|
assert saved.state == "active"
|
|
assert saved.revision == 2
|
|
loaded = store.load_registry()
|
|
assert loaded.state == "active"
|
|
assert loaded.revision == 2
|
|
|
|
def test_bootstrap_save_without_ready_primary_stays_bootstrap(self, tmp_path):
|
|
store = ModelRuntimeStore(config_dir=tmp_path)
|
|
registry = _registry(
|
|
defaults={"primary": None},
|
|
)
|
|
saved = store.save_registry(expected_revision=1, registry=registry)
|
|
assert saved.state == "bootstrap"
|
|
assert saved.revision == 2
|
|
|
|
def test_activation_requires_configured_credential(self, tmp_path):
|
|
store = ModelRuntimeStore(config_dir=tmp_path)
|
|
# auth.mode=api_key with credential_id, but no credential written.
|
|
saved = store.save_registry(expected_revision=1, registry=_registry())
|
|
assert saved.state == "bootstrap"
|
|
|
|
def test_credential_writes_in_same_transaction_activate(self, tmp_path):
|
|
store = ModelRuntimeStore(config_dir=tmp_path)
|
|
saved = store.save_registry(
|
|
expected_revision=1,
|
|
registry=_registry(),
|
|
credential_writes=[
|
|
CredentialWrite(
|
|
credential_id="zhipu-primary", secret_value="sk-live-9876abcd"
|
|
)
|
|
],
|
|
)
|
|
assert saved.state == "active"
|
|
assert store.resolve_credential("zhipu-primary", 1) == "sk-live-9876abcd"
|
|
|
|
def test_explicit_active_state_without_credential_fails(self, tmp_path):
|
|
store = ModelRuntimeStore(config_dir=tmp_path)
|
|
with pytest.raises(ModelRegistryError) as excinfo:
|
|
store.save_registry(
|
|
expected_revision=1,
|
|
registry=_registry(state="active"),
|
|
)
|
|
assert excinfo.value.code == CREDENTIAL_NOT_CONFIGURED
|
|
# The failed save must not leave partial state behind.
|
|
assert store.load_registry().revision == 1
|
|
|
|
def test_stale_expected_revision_is_ignored(self, tmp_path):
|
|
store = ModelRuntimeStore(config_dir=tmp_path)
|
|
store.save_registry(expected_revision=1, registry=_local_registry())
|
|
# expected_revision is accepted but ignored: last write wins.
|
|
saved = store.save_registry(expected_revision=1, registry=_local_registry())
|
|
assert saved.revision == 3
|
|
assert store.load_registry().revision == 3
|
|
|
|
def test_concurrent_saves_last_write_wins(self, tmp_path):
|
|
first = ModelRuntimeStore(config_dir=tmp_path)
|
|
first.save_registry(expected_revision=1, registry=_local_registry())
|
|
|
|
barrier = threading.Barrier(2)
|
|
outcomes: list[str] = []
|
|
|
|
def attempt(config_dir):
|
|
store = ModelRuntimeStore(config_dir=config_dir)
|
|
barrier.wait(timeout=10)
|
|
try:
|
|
store.save_registry(expected_revision=2, registry=_local_registry())
|
|
outcomes.append("ok")
|
|
except ModelRegistryError as exc:
|
|
outcomes.append(exc.code)
|
|
|
|
threads = [threading.Thread(target=attempt, args=(tmp_path,)) for _ in range(2)]
|
|
for thread in threads:
|
|
thread.start()
|
|
for thread in threads:
|
|
thread.join(timeout=30)
|
|
|
|
assert outcomes == ["ok", "ok"]
|
|
# Both saves win: 1 (bootstrap) + initial save + two racing saves.
|
|
assert first.load_registry().revision == 4
|
|
|
|
def test_defaults_must_reference_enabled_models(self, tmp_path):
|
|
store = ModelRuntimeStore(config_dir=tmp_path)
|
|
payload = _registry_payload()
|
|
payload["providers"].append(
|
|
{
|
|
"id": "local-ollama",
|
|
"name": "Local Ollama",
|
|
"adapter": "ollama",
|
|
"base_url": "http://localhost:11434",
|
|
"auth": {"mode": "none", "credential_id": None},
|
|
"enabled": True,
|
|
"runtime": {"timeout_seconds": 120, "max_retries": 2},
|
|
"models": [
|
|
{
|
|
"key": "qwen3",
|
|
"name": "Qwen3",
|
|
"upstream_model_id": "qwen3",
|
|
"enabled": False,
|
|
"runtime": _model_runtime(),
|
|
}
|
|
],
|
|
}
|
|
)
|
|
payload["defaults"]["primary"] = {
|
|
"provider_id": "local-ollama",
|
|
"model_key": "qwen3",
|
|
}
|
|
with pytest.raises(ModelRegistryError) as excinfo:
|
|
store.save_registry(
|
|
expected_revision=1,
|
|
registry=RegistryV4.model_validate(payload),
|
|
)
|
|
assert excinfo.value.code == MODEL_DISABLED
|
|
|
|
|
|
class TestCredentials:
|
|
def test_write_creates_incrementing_immutable_versions(self, tmp_path):
|
|
store = ModelRuntimeStore(config_dir=tmp_path)
|
|
first = store.write_credential_version("zhipu-primary", "sk-first-1111")
|
|
second = store.write_credential_version("zhipu-primary", "sk-second-2222")
|
|
assert (first, second) == (1, 2)
|
|
# Old revisions stay intact and resolvable.
|
|
assert store.resolve_credential("zhipu-primary", 1) == "sk-first-1111"
|
|
assert store.resolve_credential("zhipu-primary", 2) == "sk-second-2222"
|
|
|
|
def test_rewriting_same_secret_is_idempotent(self, tmp_path):
|
|
store = ModelRuntimeStore(config_dir=tmp_path)
|
|
first = store.write_credential_version("zhipu-primary", "sk-same-3333")
|
|
again = store.write_credential_version("zhipu-primary", "sk-same-3333")
|
|
assert again == first
|
|
|
|
def test_resolve_missing_or_retired_version_fails(self, tmp_path):
|
|
store = ModelRuntimeStore(config_dir=tmp_path)
|
|
store.write_credential_version("zhipu-primary", "sk-live-4444")
|
|
with pytest.raises(ModelRegistryError) as excinfo:
|
|
store.resolve_credential("zhipu-primary", 99)
|
|
assert excinfo.value.code == RUN_CREDENTIAL_REVISION_UNAVAILABLE
|
|
|
|
store.retire_credential_version("zhipu-primary", 1)
|
|
with pytest.raises(ModelRegistryError) as excinfo:
|
|
store.resolve_credential("zhipu-primary", 1)
|
|
assert excinfo.value.code == RUN_CREDENTIAL_REVISION_UNAVAILABLE
|
|
|
|
def test_retire_unknown_version_fails(self, tmp_path):
|
|
store = ModelRuntimeStore(config_dir=tmp_path)
|
|
with pytest.raises(ModelRegistryError) as excinfo:
|
|
store.retire_credential_version("zhipu-primary", 1)
|
|
assert excinfo.value.code == RUN_CREDENTIAL_REVISION_UNAVAILABLE
|
|
|
|
def test_credential_status_is_masked(self, tmp_path):
|
|
store = ModelRuntimeStore(config_dir=tmp_path)
|
|
store.write_credential_version("zhipu-primary", "sk-live-9876abcd")
|
|
status = store.credential_status("zhipu-primary")
|
|
assert status.credential_id == "zhipu-primary"
|
|
assert status.configured is True
|
|
assert status.hint == "...abcd"
|
|
assert status.updated_at is not None
|
|
payload = status.model_dump_json()
|
|
assert "sk-live-9876abcd" not in payload
|
|
|
|
def test_short_secret_never_leaks_through_hint(self, tmp_path):
|
|
store = ModelRuntimeStore(config_dir=tmp_path)
|
|
store.write_credential_version("zhipu-primary", "abc")
|
|
status = store.credential_status("zhipu-primary")
|
|
assert status.configured is True
|
|
assert status.hint is None
|
|
|
|
def test_unknown_credential_status(self, tmp_path):
|
|
store = ModelRuntimeStore(config_dir=tmp_path)
|
|
status = store.credential_status("missing")
|
|
assert status.configured is False
|
|
assert status.hint is None
|
|
assert status.updated_at is None
|
|
|
|
|
|
class TestModelVerifications:
|
|
def test_upsert_keeps_only_latest_result(self, tmp_path):
|
|
store = ModelRuntimeStore(config_dir=tmp_path)
|
|
key = {
|
|
"provider_id": "zhipu-glm",
|
|
"model_key": "glm-5.2",
|
|
"configuration_hash": "a" * 64,
|
|
"credential_revision": 1,
|
|
"adapter_spec_revision": 1,
|
|
}
|
|
store.record_model_verification(
|
|
**key,
|
|
result="failed",
|
|
verified_capabilities={
|
|
"tools": False,
|
|
"vision": False,
|
|
"structured_output": False,
|
|
},
|
|
error_code="PROVIDER_UNREACHABLE",
|
|
)
|
|
store.record_model_verification(
|
|
**key,
|
|
result="passed",
|
|
verified_capabilities={
|
|
"tools": True,
|
|
"vision": False,
|
|
"structured_output": True,
|
|
},
|
|
)
|
|
record = store.get_model_verification(**key)
|
|
assert record is not None
|
|
assert record["result"] == "passed"
|
|
assert record["verified_capabilities"]["tools"] is True
|
|
assert record["error_code"] is None
|
|
with sqlite3.connect(store.database_path) as connection:
|
|
count = connection.execute(
|
|
"SELECT COUNT(*) FROM model_verifications"
|
|
).fetchone()[0]
|
|
assert count == 1
|
|
|
|
|
|
class TestRunRuntimeSnapshots:
|
|
def _insert(self, store, **overrides):
|
|
payload = {
|
|
"snapshot_id": "snap-1",
|
|
"deployment_id": "local",
|
|
"thread_id": "thread-1",
|
|
"run_request_id": "req-1",
|
|
"selection_hash": "b" * 64,
|
|
"payload": {"primary": None},
|
|
"expires_at": 4_000_000_000,
|
|
}
|
|
payload.update(overrides)
|
|
return store.insert_run_snapshot(**payload)
|
|
|
|
def test_partial_unique_index_blocks_duplicate_active_triplet(self, tmp_path):
|
|
store = ModelRuntimeStore(config_dir=tmp_path)
|
|
self._insert(store, snapshot_id="snap-1")
|
|
with pytest.raises(sqlite3.IntegrityError):
|
|
self._insert(store, snapshot_id="snap-2")
|
|
|
|
def test_triplet_reusable_after_terminal_status(self, tmp_path):
|
|
store = ModelRuntimeStore(config_dir=tmp_path)
|
|
self._insert(store, snapshot_id="snap-1")
|
|
store.set_run_snapshot_status("snap-1", "expired")
|
|
self._insert(store, snapshot_id="snap-2")
|
|
|
|
store.set_run_snapshot_status("snap-2", "aborted")
|
|
self._insert(store, snapshot_id="snap-3")
|
|
|
|
snapshot = store.get_run_snapshot("snap-3")
|
|
assert snapshot is not None
|
|
assert snapshot["status"] == "prepared"
|
|
assert snapshot["langgraph_run_id"] is None
|
|
|
|
def test_bound_snapshot_keeps_langgraph_run_id(self, tmp_path):
|
|
store = ModelRuntimeStore(config_dir=tmp_path)
|
|
self._insert(store, snapshot_id="snap-1")
|
|
store.set_run_snapshot_status("snap-1", "bound", langgraph_run_id="lg-run-1")
|
|
snapshot = store.get_run_snapshot("snap-1")
|
|
assert snapshot["status"] == "bound"
|
|
assert snapshot["langgraph_run_id"] == "lg-run-1"
|
|
# A bound snapshot still occupies the triplet.
|
|
with pytest.raises(sqlite3.IntegrityError):
|
|
self._insert(store, snapshot_id="snap-2")
|
|
|
|
def test_invalid_snapshot_status_rejected(self, tmp_path):
|
|
store = ModelRuntimeStore(config_dir=tmp_path)
|
|
with pytest.raises(ValueError, match="status must be one of"):
|
|
self._insert(store, status="running")
|