Add local ML models and document suggestions (#13)

automatic-pdf-renamer/tests/test_model_manager.py

@@ -0,0 +1,124 @@
+"""Tests for full-set model retraining and shared additional-field extraction."""
+
+import json
+import os
+import tempfile
+import unittest
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import patch
+
+from app.model_manager import BASE_MODELS, list_models, train_all_models
+from app.ml import LocalMLPipeline, ModelUnavailableError
+
+
+def validated_records():
+    return [
+        SimpleNamespace(
+            extracted_text="Invoice number INV-100 for account AC-10",
+            optional_fields={"invoice_number": "INV-100", "account_number": "AC-10"},
+            emitting_company="Company A",
+            document_type="invoice",
+            template_name="invoice_template",
+            period_format="monthly",
+        ),
+        SimpleNamespace(
+            extracted_text="Invoice number INV-200",
+            optional_fields={"invoice_number": "INV-200"},
+            emitting_company="Company A",
+            document_type="invoice",
+            template_name="invoice_template",
+            period_format="quarterly",
+        ),
+        SimpleNamespace(
+            extracted_text="Account statement for AC-30",
+            optional_fields={"account_number": "AC-30"},
+            emitting_company="Company B",
+            document_type="statement",
+            template_name="statement_template",
+            period_format="yearly",
+        ),
+    ]
+
+
+class FullModelRetrainingTests(unittest.TestCase):
+    def setUp(self):
+        self.temp_dir = tempfile.TemporaryDirectory()
+        self.root = Path(self.temp_dir.name)
+        self.env = patch.dict(os.environ, {"MODEL_DIR": str(self.root)})
+        self.env.start()
+        self.addCleanup(self.env.stop)
+
+    def tearDown(self):
+        self.temp_dir.cleanup()
+
+    def test_retraining_replaces_all_models_without_history(self):
+        old_artifact = self.root / "sets" / "old-version" / "field_document_type.joblib"
+        old_artifact.parent.mkdir(parents=True)
+        old_artifact.write_bytes(b"old artifact")
+        (self.root / "active_models.json").write_text(
+            json.dumps({"field_document_type": "old-version"}), encoding="utf-8"
+        )
+
+        progress = []
+        result = train_all_models(validated_records(), lambda value, message: progress.append((value, message)))
+
+        self.assertEqual(result["status"], "completed")
+        self.assertEqual(result["count"], len(BASE_MODELS))
+        active = json.loads((self.root / "active_models.json").read_text())
+        self.assertEqual(set(active), set(BASE_MODELS))
+        self.assertEqual(len(set(active.values())), 1)
+        for name, version in active.items():
+            self.assertTrue((self.root / "sets" / version / f"{name}.joblib").is_file())
+        self.assertFalse(old_artifact.exists())
+        self.assertEqual(
+            {path.name.removesuffix(".joblib") for path in (self.root / "sets" / next(iter(active.values()))).glob("*.joblib")},
+            set(BASE_MODELS),
+        )
+        self.assertEqual(len(list((self.root / "sets").glob("*"))), 1)
+        self.assertFalse((self.root / "versions").exists())
+        self.assertEqual(progress[-1][0], 100)
+
+    def test_failed_retraining_keeps_the_previous_active_set(self):
+        old_generation = "existing-set"
+        old_artifact = self.root / "sets" / old_generation / "field_document_type.joblib"
+        old_artifact.parent.mkdir(parents=True)
+        old_artifact.write_bytes(b"previous active artifact")
+        active_before = {"field_document_type": old_generation}
+        (self.root / "active_models.json").write_text(json.dumps(active_before), encoding="utf-8")
+
+        with self.assertRaises(ValueError):
+            train_all_models([])
+
+        self.assertEqual(
+            json.loads((self.root / "active_models.json").read_text()),
+            active_before,
+        )
+        self.assertTrue(old_artifact.is_file())
+        self.assertFalse(any(path.name.startswith(".staging-") for path in self.root.iterdir()))
+
+    def test_inventory_has_no_version_history_or_emission_date_model(self):
+        result = list_models()
+        names = {model["name"] for model in result["models"]}
+        self.assertNotIn("field_emission_date", names)
+        self.assertTrue(all("versions" not in model for model in result["models"]))
+
+    def test_inference_requires_a_complete_active_model_set(self):
+        pipeline = LocalMLPipeline(str(self.root), 0.8)
+        with self.assertRaises(ModelUnavailableError):
+            pipeline._load("field_document_type")
+
+    def test_shared_additional_value_model_extracts_multiple_fields(self):
+        train_all_models(validated_records())
+        pipeline = LocalMLPipeline(str(self.root), 0.0)
+        predictions = pipeline.predict_additional_values(
+            ["invoice_number", "account_number"],
+            "Invoice number INV-200 for account AC-20",
+        )
+        self.assertEqual(set(predictions), {"invoice_number", "account_number"})
+        self.assertTrue(predictions["invoice_number"][0])
+        self.assertTrue(predictions["account_number"][0])
+
+
+if __name__ == "__main__":
+    unittest.main()