69 lines
2.2 KiB
Python
69 lines
2.2 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
|
|
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: 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
|
|
|
|
|
|
def test_registry_persists_statistical_parameters(tmp_path: Path) -> None:
|
|
from app.ml.feature_store import FeatureStore, FeatureVector
|
|
from app.ml.training import TrainingPipeline
|
|
|
|
store = FeatureStore()
|
|
store.add_batch(
|
|
[
|
|
FeatureVector("sensor.kitchen", {"temperature": 19.0}),
|
|
FeatureVector("sensor.kitchen", {"temperature": 20.0}),
|
|
]
|
|
)
|
|
artifact = TrainingPipeline(store).run("model-v1")
|
|
|
|
ModelRegistry(tmp_path).register(artifact)
|
|
|
|
assert ModelRegistry(tmp_path).load_artifact("model-v1") == artifact
|
|
|
|
|
|
def test_registry_replaces_persisted_artifact_after_restart(tmp_path: Path) -> None:
|
|
registry = ModelRegistry(tmp_path)
|
|
registry.register(TrainedArtifact("model-v1", ("sensor.kitchen",)))
|
|
replacement = TrainedArtifact("model-v1", ("sensor.bedroom",))
|
|
|
|
registry.register(replacement)
|
|
|
|
assert registry.load_artifact("model-v1") == replacement
|
|
assert ModelRegistry(tmp_path).load_artifact("model-v1") == replacement
|
|
|
|
|
|
@pytest.mark.parametrize("artifact_id", ["../escape", "nested/model", "..", ""])
|
|
def test_registry_rejects_unsafe_artifact_ids(tmp_path: 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: 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)
|