"""Upload validation/normalisation tests (app/image_guard.py).""" from __future__ import annotations import io import struct import zlib import pytest from PIL import Image, ImageFile from app.image_guard import ImageRejected, validate_and_normalize def make_png_bytes(w: int = 64, h: int = 64, color=(200, 30, 30)) -> bytes: raw = bytearray() for y in range(h): raw.append(0) for _ in range(w): raw += bytes(color) def chunk(tag, data): c = struct.pack(">I", len(data)) + tag + data return c + struct.pack(">I", zlib.crc32(tag + data) & 0xFFFFFFFF) ihdr = struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0) return (b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", ihdr) + chunk(b"IDAT", zlib.compress(bytes(raw), 9)) + chunk(b"IEND", b"")) class TestAccepts: def test_png_accepted_and_normalized(self): result = validate_and_normalize(make_png_bytes()) assert result.original_format == "PNG" assert result.width == 64 and result.height == 64 # Output really is a fresh RGB PNG. img = Image.open(io.BytesIO(result.png_bytes)) assert img.format == "PNG" and img.mode == "RGB" def test_jpeg_accepted(self): buf = io.BytesIO() Image.new("RGB", (50, 40), (10, 120, 200)).save(buf, "JPEG") result = validate_and_normalize(buf.getvalue()) assert result.original_format == "JPEG" def test_gif_accepted_first_frame(self): buf = io.BytesIO() Image.new("P", (32, 32)).save(buf, "GIF") result = validate_and_normalize(buf.getvalue()) assert result.original_format == "GIF" def test_downscale_applies(self): result = validate_and_normalize(make_png_bytes(2000, 1000), normalize_max_side=1024) assert result.downscaled is True assert max(result.width, result.height) <= 1024 def test_exif_stripped(self): img = Image.new("RGB", (24, 24), (1, 2, 3)) buf = io.BytesIO() img.save(buf, "PNG", pnginfo=None) result = validate_and_normalize(buf.getvalue()) assert b"Exif" not in result.png_bytes class TestRejects: def test_empty(self): with pytest.raises(ImageRejected) as exc: validate_and_normalize(b"") assert exc.value.code == "empty_file" def test_html_bytes(self): with pytest.raises(ImageRejected) as exc: validate_and_normalize(b"nope") assert exc.value.code in {"svg_rejected", "not_an_image"} def test_svg_rejected(self): svg = b'' with pytest.raises(ImageRejected) as exc: validate_and_normalize(svg) assert exc.value.code == "svg_rejected" def test_oversize_bytes(self): with pytest.raises(ImageRejected) as exc: validate_and_normalize(b"\x89PNG" + b"0" * 2048, max_bytes=1024) assert exc.value.code == "file_too_large" def test_truncated_png(self): with pytest.raises(ImageRejected): validate_and_normalize(make_png_bytes()[:60]) def test_pixel_cap(self): # Declared dimensions exceed the cap; Pillow raises on open. with pytest.raises(ImageRejected): validate_and_normalize(make_png_bytes(64, 64), max_pixels=100) def test_dimension_cap_precedes_pixel_decode(self, monkeypatch): decoded = False original_load = ImageFile.ImageFile.load def tracked_load(image, *args, **kwargs): nonlocal decoded decoded = True return original_load(image, *args, **kwargs) monkeypatch.setattr(ImageFile.ImageFile, "load", tracked_load) with pytest.raises(ImageRejected) as exc: validate_and_normalize( make_png_bytes(64, 64), max_pixels=1_000_000, max_side=32 ) assert exc.value.code == "image_too_large" assert decoded is False def test_random_bytes(self): with pytest.raises(ImageRejected): validate_and_normalize(b"\xde\xad\xbe\xef" * 64)