48 lines
1.6 KiB
Python
48 lines
1.6 KiB
Python
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from app.ml.feature_store import FeatureStore, FeatureVector
|
|
from app.ml.predictor import Predictor
|
|
from app.ml.training import TrainingPipeline
|
|
|
|
|
|
def _vector(sensor_id: str, temperature: float, label: str | None = None) -> FeatureVector:
|
|
return FeatureVector(sensor_id=sensor_id, values={"temperature": temperature}, label=label)
|
|
|
|
|
|
def predictor() -> Predictor:
|
|
store = FeatureStore()
|
|
store.add_batch([_vector("sensor.kitchen", 19.0), _vector("sensor.bedroom", 18.5)])
|
|
pipeline = TrainingPipeline(store)
|
|
pipeline.run("artifact_v1")
|
|
return Predictor(pipeline)
|
|
|
|
|
|
def test_predict_returns_expected_format() -> None:
|
|
p = predictor()
|
|
result = p.predict("artifact_v1", _vector("sensor.kitchen", 21.0))
|
|
assert result == "artifact_v1:sensor.kitchen:{'temperature': 21.0}"
|
|
|
|
|
|
def test_predict_rejects_unknown_sensor() -> None:
|
|
p = predictor()
|
|
with pytest.raises(ValueError):
|
|
p.predict("artifact_v1", _vector("sensor.unknown", 10.0))
|
|
|
|
|
|
def test_predict_batch_matches_single_calls() -> None:
|
|
p = predictor()
|
|
entities = [_vector("sensor.kitchen", 21.0), _vector("sensor.bedroom", 19.0)]
|
|
assert p.predict_batch("artifact_v1", entities) == [
|
|
p.predict("artifact_v1", item) for item in entities
|
|
]
|
|
|
|
|
|
def test_default_artifact_returns_last_registered() -> None:
|
|
store = FeatureStore()
|
|
store.add_batch([_vector("sensor.kitchen", 19.0), _vector("sensor.bedroom", 18.5)])
|
|
pipeline = TrainingPipeline(store)
|
|
pipeline.run("first")
|
|
pipeline.run("second")
|
|
assert Predictor.default_artifact(pipeline).artifact_id == "second" |