157 lines
7.1 KiB
Python
157 lines
7.1 KiB
Python
import io
|
|
from pathlib import Path
|
|
import struct
|
|
import sys
|
|
import unittest
|
|
from unittest.mock import patch
|
|
import zlib
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
from preview_core import HEADER, PreviewExpired, PreviewStore, encode_image, process_batch, to_pixels, validate_shape
|
|
|
|
|
|
def unpack(packet):
|
|
magic, version, format_id, mask, reserved, width, height, length, crc = HEADER.unpack_from(packet)
|
|
plain = bytes(value ^ mask for value in packet[HEADER.size:])
|
|
assert (magic, version, reserved, length, crc) == (b"CEPV", 1, 0, len(plain), zlib.crc32(plain))
|
|
return plain, (width, height), format_id
|
|
|
|
|
|
class EncodingTests(unittest.TestCase):
|
|
def test_lossless_original_and_thumbnail_dimensions(self):
|
|
rng = np.random.default_rng(9)
|
|
for shape in [(503, 777, 3), (37, 29, 4), (17, 23), (1, 1, 3)]:
|
|
with self.subTest(shape=shape):
|
|
pixels = rng.integers(0, 256, shape, dtype=np.uint8)
|
|
thumb, original, dimensions = encode_image(pixels)
|
|
plain, size, format_id = unpack(original)
|
|
self.assertEqual(size, (shape[1], shape[0]))
|
|
self.assertEqual(dimensions, {"width": shape[1], "height": shape[0]})
|
|
self.assertEqual(format_id, 1)
|
|
self.assertNotEqual(original[HEADER.size:HEADER.size + 8], b"\x89PNG\r\n\x1a\n")
|
|
with Image.open(io.BytesIO(plain)) as image:
|
|
np.testing.assert_array_equal(np.asarray(image), pixels)
|
|
thumb_plain, thumb_size, thumb_format = unpack(thumb)
|
|
self.assertLessEqual(max(thumb_size), 384)
|
|
self.assertLessEqual(abs(thumb_size[0] / thumb_size[1] - size[0] / size[1]), 0.02)
|
|
self.assertEqual(thumb_format, 1 if len(shape) == 3 and shape[2] == 4 else 2)
|
|
with Image.open(io.BytesIO(thumb_plain)) as image:
|
|
image.load()
|
|
self.assertEqual(image.size, thumb_size)
|
|
if len(shape) == 3 and shape[2] == 4:
|
|
np.testing.assert_array_equal(np.asarray(image)[..., 3], pixels[..., 3])
|
|
|
|
def test_tensor_conversion_does_not_mutate_input(self):
|
|
values = np.array([[[-0.1, 0.5, 1.1]]], dtype=np.float32)
|
|
before = values.copy()
|
|
np.testing.assert_array_equal(to_pixels(values), [[[0, 127, 255]]])
|
|
np.testing.assert_array_equal(values, before)
|
|
self.assertEqual(to_pixels(np.ones((2, 3, 1), np.float32)).shape, (2, 3))
|
|
|
|
def test_invalid_inputs(self):
|
|
for shape in [(0, 10, 10, 3), (257, 10, 10, 3), (1, 0, 10, 3), (1, 8192, 8192, 3), (1, 10, 10, 2), (10, 10, 3)]:
|
|
with self.subTest(shape=shape), self.assertRaises(ValueError):
|
|
validate_shape(shape)
|
|
for values in [np.full((1, 1, 3), np.nan), np.full((1, 1, 3), np.inf), np.ones((1, 1, 3), np.uint8)]:
|
|
with self.assertRaises(ValueError):
|
|
to_pixels(values)
|
|
|
|
|
|
class StoreTests(unittest.TestCase):
|
|
def test_expiration_and_active_batch(self):
|
|
now = [0]
|
|
store = PreviewStore(ttl=10, clock=lambda: now[0])
|
|
batch = store.begin(1)
|
|
store.add(batch, 0, b"thumb", b"original")
|
|
now[0] = 100
|
|
self.assertEqual(store.get(batch, 0, "original"), b"original")
|
|
store.finish(batch)
|
|
now[0] = 111
|
|
store.expire()
|
|
with self.assertRaises(PreviewExpired):
|
|
store.get(batch, 0, "thumbnail")
|
|
self.assertEqual((store.size, store.item_count, len(store.batches)), (0, 0, 0))
|
|
|
|
def test_eviction_keeps_current_batch_and_reports_full(self):
|
|
store = PreviewStore(max_bytes=10)
|
|
old = store.begin(1)
|
|
store.add(old, 0, b"123", b"456")
|
|
store.finish(old)
|
|
current = store.begin(2)
|
|
store.add(current, 0, b"abc", b"def")
|
|
with self.assertRaises(PreviewExpired):
|
|
store.get(old, 0, "original")
|
|
with self.assertRaisesRegex(RuntimeError, "cache is full"):
|
|
store.add(current, 1, b"123", b"456")
|
|
self.assertEqual(store.get(current, 0, "original"), b"def")
|
|
self.assertEqual(store.item_count, 1)
|
|
|
|
def test_sequence_and_entry_limits(self):
|
|
store = PreviewStore(max_items=1)
|
|
batch = store.begin(2)
|
|
with self.assertRaises(ValueError):
|
|
store.add(batch, 1, b"a", b"b")
|
|
store.add(batch, 0, b"a", b"b")
|
|
with self.assertRaises(ValueError):
|
|
store.add(batch, 0, b"a", b"b")
|
|
with self.assertRaises(RuntimeError):
|
|
store.add(batch, 1, b"a", b"b")
|
|
with self.assertRaises(ValueError):
|
|
store.get(batch, 0, "../../file")
|
|
|
|
|
|
class BatchTests(unittest.TestCase):
|
|
def test_100_images_in_order_with_small_ui_result(self):
|
|
images = np.random.default_rng(3).random((100, 8, 9, 3), dtype=np.float32)
|
|
events = []
|
|
store = PreviewStore()
|
|
result = process_batch(images, store, lambda image: image, events.append)
|
|
self.assertEqual([event["item"]["index"] for event in events if event["event"] == "image"], list(range(100)))
|
|
self.assertEqual((events[0]["event"], events[-1]["event"], result["status"]), ("start", "complete", "complete"))
|
|
for item in result["items"]:
|
|
plain, _, _ = unpack(store.get(result["batch_id"], item["index"], "original"))
|
|
with Image.open(io.BytesIO(plain)) as image:
|
|
np.testing.assert_array_equal(np.asarray(image), to_pixels(images[item["index"]]))
|
|
self.assertLess(len(str(result)), 8000)
|
|
|
|
def test_partial_failure_is_visible_and_propagated(self):
|
|
images = np.ones((3, 4, 4, 3), np.float32)
|
|
images[1, 0, 0, 0] = np.nan
|
|
events = []
|
|
store = PreviewStore()
|
|
with self.assertRaises(ValueError):
|
|
process_batch(images, store, lambda image: image, events.append)
|
|
batch_id = events[0]["batch_id"]
|
|
self.assertEqual(events[-1]["event"], "failed")
|
|
self.assertEqual(events[-1]["index"], 1)
|
|
self.assertEqual(store.item_count, 1)
|
|
self.assertTrue(store.get(batch_id, 0, "original"))
|
|
self.assertFalse(store.batches[batch_id].active)
|
|
self.assertNotIn("complete", [event["event"] for event in events])
|
|
|
|
def test_interruption_leaves_no_unpublished_image(self):
|
|
calls = [0]
|
|
def interrupt():
|
|
calls[0] += 1
|
|
if calls[0] == 2:
|
|
raise InterruptedError("cancelled")
|
|
store, events = PreviewStore(), []
|
|
with self.assertRaises(InterruptedError):
|
|
process_batch(np.ones((1, 4, 4, 3), np.float32), store, lambda image: image, events.append, interrupt)
|
|
self.assertEqual(store.item_count, 0)
|
|
self.assertEqual(events[-1]["event"], "failed")
|
|
|
|
def test_transport_error_does_not_hide_original_error(self):
|
|
def notify(event):
|
|
if event["event"] == "failed":
|
|
raise RuntimeError("transport unavailable")
|
|
with self.assertRaises(ValueError):
|
|
process_batch(np.full((1, 2, 2, 3), np.nan), PreviewStore(), lambda image: image, notify)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|