"""Setup integrity checks. All response bytes are local fakes; no network.""" import hashlib import importlib.util import io import pathlib import tempfile import unittest spec = importlib.util.spec_from_file_location('download_model', pathlib.Path(__file__).parents[1] / 'scripts/download_model.py') module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) class Response(io.BytesIO): headers = {} class SetupIntegrity(unittest.TestCase): def setUp(self): self.temp = tempfile.TemporaryDirectory() self.addCleanup(self.temp.cleanup) self.path = pathlib.Path(self.temp.name) / 'artifact.bin' self.good = b'fixed synthetic model bytes' self.sha = hashlib.sha256(self.good).hexdigest() def read(self, opener): return module.read_verified(self.path, 'https://example.invalid/pinned', len(self.good), self.sha, opener) def test_existing_valid_cache_never_downloads(self): self.path.write_bytes(self.good) def forbidden(*args, **kwargs): raise AssertionError('Unexpected network request') self.assertEqual(self.read(forbidden), self.good) def test_corrupt_cache_is_rejected_and_preserved(self): self.path.write_bytes(b'corrupt') with self.assertRaisesRegex(ValueError, 'integrity mismatch'): self.read(lambda *args, **kwargs: Response(self.good)) self.assertEqual(self.path.read_bytes(), b'corrupt') def test_verified_download_saved_atomically(self): self.assertEqual(self.read(lambda *args, **kwargs: Response(self.good)), self.good) self.assertEqual(self.path.read_bytes(), self.good) self.assertFalse(self.path.with_suffix('.bin.part').exists()) def test_untrusted_download_not_saved(self): with self.assertRaisesRegex(ValueError, 'integrity mismatch'): self.read(lambda *args, **kwargs: Response(b'wrong bytes')) self.assertFalse(self.path.exists()) if __name__ == '__main__': unittest.main()