tests/test_integrity.py

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