cesnimda
41 lines
1.8 KiB
Python
41 lines
1.8 KiB
Python
from pathlib import Path
|
|
import copy
|
|
import hashlib
|
|
import tempfile
|
|
import unittest
|
|
from verify import verify_fixture
|
|
|
|
|
|
class ArtifactIntegrity(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.temp.cleanup)
|
|
self.directory = Path(self.temp.name)
|
|
(self.directory / 'model.safetensors').write_bytes(b'known fixture')
|
|
self.manifest = {'files': [{'file': 'model.safetensors', 'bytes': 13, 'sha256': hashlib.sha256(b'known fixture').hexdigest()}]}
|
|
|
|
def test_accepts_exact_known_content(self):
|
|
self.assertEqual(verify_fixture(self.directory, self.manifest), self.directory.resolve())
|
|
|
|
def test_rejects_same_size_corruption(self):
|
|
(self.directory / 'model.safetensors').write_bytes(b'wrong fixture')
|
|
with self.assertRaises(ValueError): verify_fixture(self.directory, self.manifest)
|
|
|
|
def test_rejects_missing_weight(self):
|
|
(self.directory / 'model.safetensors').unlink()
|
|
with self.assertRaises(ValueError): verify_fixture(self.directory, self.manifest)
|
|
|
|
def test_rejects_traversal_and_duplicate_identity(self):
|
|
for value in ['../model.safetensors', '/model.safetensors']:
|
|
manifest = copy.deepcopy(self.manifest); manifest['files'][0]['file'] = value
|
|
with self.assertRaises(ValueError): verify_fixture(self.directory, manifest)
|
|
self.manifest['files'] *= 2
|
|
with self.assertRaises(ValueError): verify_fixture(self.directory, self.manifest)
|
|
|
|
def test_rejects_unlisted_model_code(self):
|
|
(self.directory / 'modeling_unreviewed.py').write_text('raise RuntimeError()')
|
|
with self.assertRaises(ValueError): verify_fixture(self.directory, self.manifest)
|
|
|
|
|
|
if __name__ == '__main__': unittest.main()
|