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