prototype/tests/setup_test.py

"""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()