39 lines
1.2 KiB
Python
39 lines
1.2 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
import pytest
|
|
|
|
from app.ml.registry.model_registry import ModelRegistry
|
|
from app.ml.training import TrainedArtifact
|
|
|
|
|
|
def test_registry_loads_persisted_artifacts_after_restart(tmp_path) -> None:
|
|
registry = ModelRegistry(tmp_path)
|
|
artifact = TrainedArtifact("model-v1", ("sensor.kitchen", "sensor.bedroom"))
|
|
registry.register(artifact)
|
|
|
|
restarted = ModelRegistry(tmp_path)
|
|
|
|
assert restarted.load_artifact("model-v1") == artifact
|
|
|
|
|
|
@pytest.mark.parametrize("artifact_id", ["../escape", "nested/model", "..", ""])
|
|
def test_registry_rejects_unsafe_artifact_ids(tmp_path, artifact_id: str) -> None:
|
|
registry = ModelRegistry(tmp_path)
|
|
|
|
with pytest.raises(ValueError):
|
|
registry.register(TrainedArtifact(artifact_id, ("sensor.kitchen",)))
|
|
|
|
assert list(tmp_path.parent.glob("escape.json")) == []
|
|
|
|
|
|
def test_registry_rejects_corrupt_persisted_artifact(tmp_path) -> None:
|
|
(tmp_path / "broken.json").write_text(
|
|
json.dumps({"artifact_id": "../broken", "supported_sensors": []}),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="broken.json"):
|
|
ModelRegistry(tmp_path)
|