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