Spaces:
Running
Running
| """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"<html><body>nope</body></html>") | |
| assert exc.value.code in {"svg_rejected", "not_an_image"} | |
| def test_svg_rejected(self): | |
| svg = b'<?xml version="1.0"?><svg xmlns="http://www.w3.org/2000/svg"><script>alert(1)</script></svg>' | |
| 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) | |