Verify cached model checksum before use

tests/test_aiprocess.py

@@ -1,16 +1,21 @@
 import sys
+import tempfile
 import unittest
+from hashlib import sha256
 from pathlib import Path
+from unittest.mock import patch
 
 ROOT = Path(__file__).resolve().parents[1]
 SRC = ROOT / "src"
 if str(SRC) not in sys.path:
     sys.path.insert(0, str(SRC))
 
+import aiprocess
 from aiprocess import (
     _assert_expected_model_checksum,
     _extract_cleaned_text,
     _supports_response_format,
+    ensure_model,
 )
 from constants import MODEL_SHA256
 
@@ -98,5 +103,81 @@ class ModelChecksumTests(unittest.TestCase):
             _assert_expected_model_checksum("0" * 64)
 
 
+class _Response:
+    def __init__(self, payload: bytes):
+        self.payload = payload
+        self.offset = 0
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        return False
+
+    def getheader(self, name: str):
+        if name.lower() == "content-length":
+            return str(len(self.payload))
+        return None
+
+    def read(self, size: int) -> bytes:
+        if self.offset >= len(self.payload):
+            return b""
+        chunk = self.payload[self.offset : self.offset + size]
+        self.offset += len(chunk)
+        return chunk
+
+
+class EnsureModelTests(unittest.TestCase):
+    def test_existing_valid_model_skips_download(self):
+        payload = b"valid-model"
+        checksum = sha256(payload).hexdigest()
+        with tempfile.TemporaryDirectory() as td:
+            model_path = Path(td) / "model.gguf"
+            model_path.write_bytes(payload)
+            with patch.object(aiprocess, "MODEL_PATH", model_path), patch.object(
+                aiprocess, "MODEL_DIR", model_path.parent
+            ), patch.object(aiprocess, "MODEL_SHA256", checksum), patch(
+                "aiprocess.urllib.request.urlopen"
+            ) as urlopen:
+                out = ensure_model()
+
+        self.assertEqual(out, model_path)
+        urlopen.assert_not_called()
+
+    def test_existing_invalid_model_triggers_redownload(self):
+        cached_payload = b"bad-model"
+        downloaded_payload = b"good-model"
+        expected_checksum = sha256(downloaded_payload).hexdigest()
+        with tempfile.TemporaryDirectory() as td:
+            model_path = Path(td) / "model.gguf"
+            model_path.write_bytes(cached_payload)
+            with patch.object(aiprocess, "MODEL_PATH", model_path), patch.object(
+                aiprocess, "MODEL_DIR", model_path.parent
+            ), patch.object(aiprocess, "MODEL_SHA256", expected_checksum), patch(
+                "aiprocess.urllib.request.urlopen",
+                return_value=_Response(downloaded_payload),
+            ) as urlopen:
+                out = ensure_model()
+
+            self.assertEqual(out, model_path)
+            self.assertEqual(model_path.read_bytes(), downloaded_payload)
+            urlopen.assert_called_once()
+
+    def test_invalid_cached_model_and_redownload_failure_raises_clear_error(self):
+        with tempfile.TemporaryDirectory() as td:
+            model_path = Path(td) / "model.gguf"
+            model_path.write_bytes(b"bad-model")
+            with patch.object(aiprocess, "MODEL_PATH", model_path), patch.object(
+                aiprocess, "MODEL_DIR", model_path.parent
+            ), patch.object(aiprocess, "MODEL_SHA256", "f" * 64), patch(
+                "aiprocess.urllib.request.urlopen",
+                side_effect=RuntimeError("network down"),
+            ):
+                with self.assertRaisesRegex(
+                    RuntimeError, "cached model checksum mismatch and redownload failed"
+                ):
+                    ensure_model()
+
+
 if __name__ == "__main__":
     unittest.main()