File size: 5,121 Bytes
cd07add 3be79c6 cd07add 3be79c6 6c516b6 3be79c6 6c516b6 3be79c6 6c516b6 3be79c6 6c516b6 3be79c6 cd07add 3be79c6 6c516b6 3be79c6 cd07add 6c516b6 cd07add 6c516b6 cd07add 6c516b6 cd07add 6c516b6 cd07add 6c516b6 cd07add 6c516b6 cd07add 7658255 cd07add 7658255 cd07add 7658255 cd07add 7658255 cd07add 7658255 cd07add 7658255 cd07add | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | import unittest
from unittest.mock import MagicMock, patch
from app.ml.runtime import EmailRiskModel, get_model
class RuntimeLoaderTests(unittest.TestCase):
"""The ML layer must be strictly optional: any failure degrades to "no ML"
and validation continues deterministically.
Backend = PyTorch transformers, NOT ONNX: the ONNX export produced
confidently-wrong labels under two transformers versions while the PyTorch
checkpoint scored 0.999 held-out (see VALIDATOR-PERFORMANCE-REPORT).
"""
def test_load_uses_the_injected_factory(self):
fake_pipe = MagicMock()
calls = {}
def factory(repo, batch_size):
calls["args"] = (repo, batch_size)
return fake_pipe
m = EmailRiskModel(repo="user/model",
batch_size=32, pipeline_factory=factory)
self.assertTrue(m.load())
self.assertTrue(m.loaded)
self.assertEqual(calls["args"], ("user/model", 32))
self.assertIs(m._pipeline, fake_pipe)
def test_factory_failure_returns_false_and_never_raises(self):
def bad_factory(repo, batch_size):
raise RuntimeError("hub down")
m = EmailRiskModel(repo="user/model", pipeline_factory=bad_factory)
self.assertFalse(m.load())
self.assertFalse(m.loaded)
self.assertIsNone(m._pipeline)
def test_load_uses_pytorch_transformers_by_default(self):
from app.ml.runtime import default_pipeline_factory
m = EmailRiskModel(repo="user/model")
self.assertIs(m._factory, default_pipeline_factory)
def test_predict_before_load_returns_none(self):
m = EmailRiskModel(repo="user/model")
self.assertIsNone(m.predict(["a [SEP] b.test"]))
def test_predict_maps_label_and_confidence(self):
m = EmailRiskModel(repo="user/model")
fake = MagicMock(return_value=[{"label": "DISPOSABLE", "score": 0.93}])
m._pipeline = fake
m._loaded = True
out = m.predict(["x [SEP] tempmail.example"])
self.assertEqual(out, [("DISPOSABLE", 0.93)])
fake.assert_called_once()
self.assertIn("batch_size", fake.call_args.kwargs)
def test_predict_error_returns_none_not_raises(self):
m = EmailRiskModel(repo="user/model")
fake = MagicMock(side_effect=RuntimeError("onnx exploded"))
m._pipeline = fake
m._loaded = True
self.assertIsNone(m.predict(["a [SEP] b.test"]))
def test_predict_normalizes_label_case(self):
m = EmailRiskModel(repo="user/model")
fake = MagicMock(return_value=[{"label": "disposable", "score": 0.5}])
m._pipeline = fake
m._loaded = True
self.assertEqual(m.predict(["a [SEP] b.test"])[0][0], "DISPOSABLE")
def test_predict_accepts_plain_string_batch_return(self):
m = EmailRiskModel(repo="user/model")
fake = MagicMock(return_value=[[{"label": "NORMAL", "score": 0.7}]])
m._pipeline = fake
m._loaded = True
self.assertEqual(m.predict(["a [SEP] b.test"]), [("NORMAL", 0.7)])
def test_no_token_attribute_exists(self):
# Part B Change 5: the deployment surface carries NO HF token. The
# advisory model loads from the public Hub without credentials.
m = EmailRiskModel(repo="user/model")
self.assertFalse(hasattr(m, "token"))
class GetModelTests(unittest.TestCase):
def setUp(self):
from app.ml import runtime
self.runtime = runtime
runtime._model = None # reset the process-wide singleton
def tearDown(self):
self.runtime._model = None
def _patch_settings(self, enabled, repo):
return (
patch.object(self.runtime.settings, "ml_disposable_enabled", enabled),
patch.object(self.runtime.settings, "ml_model_repo", repo),
)
def test_returns_none_when_disabled(self):
p1, p2 = self._patch_settings(False, "user/model")
with p1, p2:
self.assertIsNone(get_model())
def test_returns_none_when_no_repo(self):
p1, p2 = self._patch_settings(True, "")
with p1, p2:
self.assertIsNone(get_model())
def test_returns_model_when_enabled_and_configured(self):
p1, p2 = self._patch_settings(True, "user/model")
with p1, p2, patch.object(
EmailRiskModel, "load", return_value=True
) as load_mock:
m = get_model()
self.assertIsInstance(m, EmailRiskModel)
# The model must actually be LOADED - an unloaded model fail-opens to
# None on every predict, silently disabling the feature (found live).
self.assertTrue(load_mock.called)
def test_load_failure_still_returns_model_but_fail_open(self):
p1, p2 = self._patch_settings(True, "user/model")
with p1, p2, patch.object(EmailRiskModel, "load", return_value=False):
m = get_model()
self.assertIsInstance(m, EmailRiskModel)
self.assertFalse(m.loaded)
self.assertIsNone(m.predict(["a [SEP] b.test"]))
if __name__ == "__main__":
unittest.main() |