31 lines
998 B
Python
31 lines
998 B
Python
import io
|
|
import unittest
|
|
|
|
from PIL import Image
|
|
|
|
from app.vision.preprocess import prepare_image
|
|
|
|
|
|
class VisionPreprocessTests(unittest.TestCase):
|
|
def _make_png(self, width: int, height: int) -> bytes:
|
|
buffer = io.BytesIO()
|
|
Image.new("RGB", (width, height), color=(120, 80, 200)).save(buffer, format="PNG")
|
|
return buffer.getvalue()
|
|
|
|
def test_resize_large_image(self) -> None:
|
|
raw = self._make_png(2400, 1600)
|
|
prepared = prepare_image(raw)
|
|
self.assertLessEqual(max(prepared.width, prepared.height), 1280)
|
|
self.assertLess(prepared.compressed_bytes, prepared.original_bytes)
|
|
|
|
def test_small_image_keeps_dimensions(self) -> None:
|
|
raw = self._make_png(640, 480)
|
|
prepared = prepare_image(raw)
|
|
self.assertEqual(prepared.width, 640)
|
|
self.assertEqual(prepared.height, 480)
|
|
self.assertEqual(prepared.mime, "image/jpeg")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|