71 lines
2.3 KiB
Python
71 lines
2.3 KiB
Python
from __future__ import annotations
|
|
|
|
import unittest
|
|
|
|
import pandas as pd
|
|
|
|
from engin.model import PredictionBundle, PrecomputedPredictionModel
|
|
from engin.service import DiagnosticService
|
|
from engin.validation import ValidationResult
|
|
|
|
|
|
class FakeReader:
|
|
def __init__(self, frame: pd.DataFrame) -> None:
|
|
self.frame = frame
|
|
|
|
def read_bytes(self, payload: bytes) -> pd.DataFrame:
|
|
if payload != b"valid":
|
|
raise AssertionError("unexpected payload")
|
|
return self.frame.copy()
|
|
|
|
def read_path(self, path):
|
|
return self.frame.copy()
|
|
|
|
|
|
class FakeValidator:
|
|
def __init__(self) -> None:
|
|
self.calls = 0
|
|
|
|
def validate(self, frame: pd.DataFrame) -> ValidationResult:
|
|
self.calls += 1
|
|
return ValidationResult(frame.copy(), ("warning",))
|
|
|
|
|
|
class FakeModel:
|
|
def __init__(self) -> None:
|
|
self.calls = 0
|
|
|
|
def predict(self, frame: pd.DataFrame) -> PredictionBundle:
|
|
self.calls += 1
|
|
submission = frame[["engine_id", "cylinder"]].copy()
|
|
submission["label"] = "ok"
|
|
submission["severity"] = "nie_dotyczy"
|
|
diagnostics = submission.copy()
|
|
diagnostics["label_confidence"] = 0.99
|
|
return PredictionBundle(submission, diagnostics)
|
|
|
|
|
|
class DiagnosticServiceTests(unittest.TestCase):
|
|
def test_dependencies_are_called_once_and_result_is_composed(self) -> None:
|
|
frame = pd.DataFrame({"engine_id": ["e1"], "cylinder": [1]})
|
|
validator = FakeValidator()
|
|
model = FakeModel()
|
|
service = DiagnosticService(
|
|
reader=FakeReader(frame), validator=validator, model=model
|
|
)
|
|
result = service.diagnose_bytes(b"valid")
|
|
self.assertEqual(validator.calls, 1)
|
|
self.assertEqual(model.calls, 1)
|
|
self.assertEqual(result.engine_ids, ["e1"])
|
|
self.assertEqual(result.warnings, ("warning",))
|
|
|
|
def test_precomputed_model_rejects_different_key_set(self) -> None:
|
|
submission = pd.DataFrame(
|
|
{"engine_id": ["e1"], "cylinder": [1], "label": ["ok"], "severity": ["nie_dotyczy"]}
|
|
)
|
|
diagnostics = submission.copy()
|
|
model = PrecomputedPredictionModel(submission, diagnostics)
|
|
mismatched = pd.DataFrame({"engine_id": ["e2"], "cylinder": [1]})
|
|
with self.assertRaisesRegex(Exception, "awaryjny"):
|
|
model.predict(mismatched)
|