56 lines
2.3 KiB
Python
56 lines
2.3 KiB
Python
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()
|