import tempfile import unittest from pathlib import Path from PIL import Image from split_nsfw_test_image import split_image class SplitImageTests(unittest.TestCase): def setUp(self): self.temporary = tempfile.TemporaryDirectory() self.addCleanup(self.temporary.cleanup) self.root = Path(self.temporary.name) self.source = self.root / "source.png" self.image = Image.new("RGBA", (4, 3)) self.image.putdata([(i * 21, 255 - i * 21, i * 7, i * 23) for i in range(12)]) self.image.save(self.source) def test_random_shares_losslessly_reconstruct_pixels_and_alpha(self): first, second = split_image(self.source, self.root / "one") self.assertEqual((first.name, second.name), ("image_1a.png", "image_1b.png")) with Image.open(first) as a, Image.open(second) as b: self.assertEqual(a.mode, "RGBA") self.assertEqual(a.size, self.image.size) self.assertEqual(b.size, self.image.size) reconstructed = bytes(x ^ y for x, y in zip(a.tobytes(), b.tobytes())) self.assertEqual(reconstructed, self.image.tobytes()) other_first, _ = split_image(self.source, self.root / "two") self.assertNotEqual(first.read_bytes(), other_first.read_bytes()) def test_refuses_to_overwrite_and_supports_numbered_pairs(self): directory = self.root / "fixtures" first, second = split_image(self.source, directory) original = (first.read_bytes(), second.read_bytes()) with self.assertRaises(FileExistsError): split_image(self.source, directory) self.assertEqual((first.read_bytes(), second.read_bytes()), original) next_pair = split_image(self.source, directory, index=2) self.assertEqual(tuple(path.name for path in next_pair), ("image_2a.png", "image_2b.png")) def test_rejects_nonpositive_indices_and_nonimages_without_writing_parts(self): directory = self.root / "fixtures" for index in (0, -1): with self.assertRaises(ValueError): split_image(self.source, directory, index) invalid = self.root / "invalid.txt" invalid.write_text("not an image") with self.assertRaises(OSError): split_image(invalid, directory) self.assertFalse(directory.exists()) if __name__ == "__main__": unittest.main()