File size: 1,739 Bytes
6f8c80a | 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 41 42 43 44 45 | 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 d:
p = Path(d) / 'weight.safetensors'
p.write_bytes(b'bad')
rows = [{'path': p.name, 'size': 3,
'sha256': hashlib.sha256(b'yes').hexdigest()}]
with self.assertRaisesRegex(ValueError, 'SHA256'):
verify_files(Path(d), rows)
p.write_bytes(b'yes')
self.assertEqual(len(verify_files(Path(d), rows)), 1)
def test_missing_file_cannot_be_counted_as_verified(self):
with tempfile.TemporaryDirectory() as d:
with self.assertRaises(FileNotFoundError):
verify_files(Path(d), [{'path': 'missing', 'size': 3, 'sha256': 'x'}])
def test_pair_rejects_different_edit_input_or_generation_setting(self):
row = {'case_id': 'edit', 'seed': 42, 'prompt': 'red shirt',
'width': 1024, 'height': 1024, 'steps': 40,
'cfg': 1.0, 'offload': 'model', 'input_sha256': 'abc',
'kv_cache': True}
validate_pair(row, copy.deepcopy(row))
for key, value in [('input_sha256', 'def'), ('steps', 20), ('seed', 9)]:
changed = dict(row, **{key: value})
with self.subTest(key=key), self.assertRaises(ValueError):
validate_pair(row, changed)
def test_incomplete_pair_is_rejected(self):
with self.assertRaises(ValueError):
validate_pair({}, {})
if __name__ == '__main__':
unittest.main()
|