File size: 1,686 Bytes
9116984 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 | import copy
import hashlib
import tempfile
import unittest
from pathlib import Path
from scripts.integrity import validate_pair, verify_files
class IntegrityTests(unittest.TestCase):
def test_corrupted_weight_is_rejected_even_if_size_matches(self):
with tempfile.TemporaryDirectory() as directory:
path = Path(directory) / 'weight.safetensors'
path.write_bytes(b'bad')
rows = [{'path': path.name, 'size': 3, 'sha256': hashlib.sha256(b'yes').hexdigest()}]
with self.assertRaisesRegex(ValueError, 'SHA256'):
verify_files(Path(directory), rows)
path.write_bytes(b'yes')
self.assertEqual(len(verify_files(Path(directory), rows)), 1)
def test_missing_file_cannot_be_counted_as_verified(self):
with tempfile.TemporaryDirectory() as directory:
with self.assertRaises(FileNotFoundError):
verify_files(Path(directory), [{'path': 'missing', 'size': 3, 'sha256': 'x'}])
def test_pair_rejects_a_changed_edit_seed_or_setting(self):
row = {'case_id': 'edit', 'seed': 1000042, 'prompt': 'red sweater',
'width': 1024, 'height': 1024, 'steps': 40,
'cfg': 1.0, 'offload': 'model', 'input_sha256': 'abc',
'kv_cache': True, 'source_seed': 42}
validate_pair(row, copy.deepcopy(row))
for key, value in [('input_sha256', 'def'), ('steps', 20), ('seed', 42), ('source_seed', 7)]:
changed = dict(row, **{key: value})
with self.subTest(key=key), self.assertRaises(ValueError):
validate_pair(row, changed)
if __name__ == '__main__':
unittest.main()
|