Z User
v3.1.0: Part B architecture migration — eliminate Hugging Face inference credit dependency
6c516b6 | 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() |