tests/test_editing_protocol.py

import tempfile
import unittest
from pathlib import Path

import numpy as np
from PIL import Image
from scripts import editing_protocol as protocol


class EditingProtocolTests(unittest.TestCase):
    def test_rejects_source_seed_collision_even_at_different_resolution(self):
        with self.assertRaisesRegex(ValueError, 'source seed'):
            protocol.validate_seeds([{'source_seed': 42}], [42, 123])

    def test_rejects_duplicate_empty_and_out_of_range_seeds(self):
        for seeds in ([], [11, 11], [-1], [2**63]):
            with self.subTest(seeds=seeds), self.assertRaises(ValueError):
                protocol.validate_seeds([{'source_seed': 42}], seeds)

    def test_independent_seeds_accepted_without_remapping(self):
        self.assertEqual(protocol.validate_seeds([{'source_seed': 42}], [1000042, 1000123]),
                         [1000042, 1000123])

    def test_palette_transparency_and_exif_are_respected(self):
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory)/'palette.png'
            im = Image.new('P', (2, 1), 0)
            im.putpalette([255, 0, 0] + [0, 0, 0]*255)
            im.info['transparency'] = 0
            im.save(path)
            loaded = protocol.load_input(path)
            self.assertEqual(loaded.mode, 'RGBA')
            self.assertEqual(loaded.getpixel((0, 0)), (255, 0, 0, 0))
            exif_path = Path(directory)/'rotated.png'
            im = Image.new('RGB', (32, 64))
            exif = im.getexif(); exif[274] = 6
            im.save(exif_path, exif=exif)
            self.assertEqual(protocol.load_input(exif_path).size, (64, 32))

    def test_geometry_preserves_aspect_and_uses_area_not_width(self):
        self.assertEqual(protocol.edit_dimensions((1600, 900), 1024), (1376, 768))
        self.assertEqual(protocol.edit_dimensions((900, 1600), 1024), (768, 1376))
        with self.assertRaises(ValueError):
            protocol.edit_dimensions((10, 10), 0)

    def test_alpha_diagnostics_distinguish_white_from_transparency(self):
        white = protocol.image_diagnostics(Image.new('RGBA', (2, 2), 'white'))
        self.assertEqual(white['alpha_transparent_fraction'], 0)
        im = Image.fromarray(np.array([[[255, 0, 0, 0], [0, 255, 0, 128]]], dtype=np.uint8))
        stats = protocol.image_diagnostics(im)
        self.assertEqual(stats['alpha_transparent_fraction'], .5)
        self.assertEqual(stats['alpha_soft_fraction'], .5)
        self.assertEqual(stats['alpha_min'], 0)
        before = im.tobytes()
        with tempfile.TemporaryDirectory() as directory:
            protocol.save_diagnostics(im, Path(directory), 'sample')
            black = Image.open(Path(directory)/'sample-black.png')
            self.assertEqual(black.getpixel((0, 0)), (0, 0, 0))
            self.assertEqual(black.getpixel((1, 0)), (0, 128, 0))
            self.assertEqual(im.tobytes(), before)


if __name__ == '__main__':
    unittest.main()