gendesign/tradein-mvp/backend/tests/services/test_image_sanitizer.py

"""Unit tests for `app.services.image_sanitizer.sanitize_image`.

Coverage:
- happy path: a small RGBA PNG re-encodes to a clean JPEG (RGB, EXIF stripped)
- resize path: an image above _MAX_DIMENSION is thumbnailed down, still decodes
- pixel-flood DoS: an image declaring > _MAX_PIXELS is rejected via img.size
  BEFORE img.load() is ever called (no full decode → no OOM) — #2214
- byte-size backstop: content over _MAX_BYTES rejected before Image.open
- Pillow's own bomb threshold (MAX_IMAGE_PIXELS) is lowered to our cap
- garbage bytes still raise ImageSanitizationError (regression guard)

The oversized-dimension case is asserted via monkeypatch rather than a real
100 MP buffer: allocating the bomb is exactly the OOM we are defending against,
so the test proves img.load() is NOT reached.
"""

from __future__ import annotations

import io
import os
from unittest.mock import MagicMock

import pytest

# DATABASE_URL required by config before any app import (см. sibling service tests).
os.environ.setdefault("DATABASE_URL", "postgresql+psycopg://test:test@localhost:5432/test")

from PIL import Image

from app.services import image_sanitizer
from app.services.image_sanitizer import (
    SANITIZED_CONTENT_TYPE,
    ImageSanitizationError,
    sanitize_image,
)


def _png_bytes(width: int, height: int, mode: str = "RGBA") -> bytes:
    buf = io.BytesIO()
    Image.new(mode, (width, height), color=(120, 60, 30, 255)[: len(mode)]).save(buf, format="PNG")
    return buf.getvalue()


def test_small_image_passes_and_becomes_jpeg() -> None:
    out, ctype = sanitize_image(_png_bytes(100, 80))
    assert ctype == SANITIZED_CONTENT_TYPE
    # Output is a valid JPEG, decoded back in RGB (transparency dropped).
    reloaded = Image.open(io.BytesIO(out))
    assert reloaded.format == "JPEG"
    assert reloaded.mode == "RGB"
    assert reloaded.size == (100, 80)


def test_oversized_but_under_pixel_cap_is_resized() -> None:
    # 3000x2000 = 6 MP (< 40 MP cap) but longest edge > _MAX_DIMENSION → thumbnailed.
    out, _ = sanitize_image(_png_bytes(3000, 2000, mode="RGB"))
    reloaded = Image.open(io.BytesIO(out))
    assert max(reloaded.size) == image_sanitizer._MAX_DIMENSION
    assert reloaded.size == (2400, 1600)


def test_pixel_flood_rejected_without_decode(monkeypatch) -> None:
    """Image declaring > cap pixels is rejected via img.size, never img.load()."""
    fake_img = MagicMock()
    fake_img.size = (10_000, 10_000)  # 100 MP >> 40 MP cap
    fake_img.__enter__.return_value = fake_img
    fake_img.__exit__.return_value = False
    monkeypatch.setattr(image_sanitizer.Image, "open", lambda *a, **k: fake_img)

    with pytest.raises(ImageSanitizationError) as excinfo:
        sanitize_image(b"\x89PNG\r\n\x1a\n-tiny-header")

    assert "exceeds" in str(excinfo.value)
    # The whole point of the fix: no full decode happened → no hundreds-of-MB buffer.
    fake_img.load.assert_not_called()


def test_pixel_cap_boundary_allows_at_cap(monkeypatch) -> None:
    """Exactly _MAX_PIXELS is allowed through to load() (strict '>' comparison)."""
    fake_img = MagicMock()
    fake_img.size = (image_sanitizer._MAX_PIXELS, 1)  # width*height == cap
    fake_img.mode = "RGB"
    fake_img.__enter__.return_value = fake_img
    fake_img.__exit__.return_value = False
    monkeypatch.setattr(image_sanitizer.Image, "open", lambda *a, **k: fake_img)
    monkeypatch.setattr(image_sanitizer.ImageOps, "exif_transpose", lambda im: im)

    sanitize_image(b"whatever")
    fake_img.load.assert_called_once()


def test_byte_size_backstop_rejects_before_open(monkeypatch) -> None:
    called = {"open": False}

    def _guard(*_a, **_k):
        called["open"] = True
        raise AssertionError("Image.open must not be called for oversized bytes")

    monkeypatch.setattr(image_sanitizer.Image, "open", _guard)
    with pytest.raises(ImageSanitizationError):
        sanitize_image(b"\x00" * (image_sanitizer._MAX_BYTES + 1))
    assert called["open"] is False


def test_pillow_bomb_threshold_lowered_to_cap() -> None:
    assert image_sanitizer._MAX_PIXELS == 40_000_000
    assert Image.MAX_IMAGE_PIXELS == image_sanitizer._MAX_PIXELS


def test_garbage_bytes_rejected() -> None:
    with pytest.raises(ImageSanitizationError):
        sanitize_image(b"definitely not an image")