ML-002: Trainingspipeline mit Tainted-Data-Check
This commit is contained in:
48
tests/ml/test_training.py
Normal file
48
tests/ml/test_training.py
Normal file
@@ -0,0 +1,48 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from app.ml.feature_store import FeatureStore, FeatureVector
|
||||
from app.ml.training import TrainingPipeline, TrainedArtifact
|
||||
|
||||
|
||||
def _vector(sensor_id: str, temperature: float, label: str | None = None) -> FeatureVector:
|
||||
return FeatureVector(sensor_id=sensor_id, values={"temperature": temperature}, label=label)
|
||||
|
||||
|
||||
def store_with_data() -> TrainingPipeline:
|
||||
store = FeatureStore()
|
||||
store.add_batch(
|
||||
[
|
||||
_vector("sensor.kitchen", 19.0),
|
||||
_vector("sensor.kitchen", 20.0),
|
||||
_vector("sensor.bedroom", 18.5),
|
||||
]
|
||||
)
|
||||
return TrainingPipeline(store)
|
||||
|
||||
|
||||
def test_run_returns_trained_artifact() -> None:
|
||||
pipeline = store_with_data()
|
||||
artifact = pipeline.run("artifact_v1")
|
||||
assert artifact.artifact_id == "artifact_v1"
|
||||
assert artifact.supported_sensors == ("sensor.bedroom", "sensor.kitchen")
|
||||
|
||||
|
||||
def test_run_without_data_raises_value_error() -> None:
|
||||
pipeline = TrainingPipeline(FeatureStore())
|
||||
with pytest.raises(ValueError):
|
||||
pipeline.run("artifact_v1")
|
||||
|
||||
|
||||
def test_export_returns_registered_artifact() -> None:
|
||||
pipeline = store_with_data()
|
||||
pipeline.run("artifact_v1")
|
||||
exported = pipeline.export("artifact_v1")
|
||||
assert exported == pipeline.export("artifact_v1")
|
||||
|
||||
|
||||
def test_export_missing_artifact_raises_key_error() -> None:
|
||||
pipeline = store_with_data()
|
||||
with pytest.raises(KeyError):
|
||||
pipeline.export("artifact_v1")
|
||||
Reference in New Issue
Block a user