ImpTune/tests/test_icon_upload.py

"""Integration tests for icon upload endpoint — PKG-04."""
from __future__ import annotations

import io

import pytest
from PIL import Image


def _make_png(width: int = 256, height: int = 256, size_bytes: int | None = None) -> bytes:
    """Create a valid PNG image with given dimensions in memory."""
    img = Image.new("RGBA", (width, height), color="red")
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    data = buf.getvalue()
    if size_bytes is not None and size_bytes > len(data):
        # Pad the PNG by embedding extra data (won't affect PIL open, but will exceed size limit)
        # Instead, use a raw bytes approach: return oversized raw content
        data = data + b"\x00" * (size_bytes - len(data))
    return data


def _make_jpeg(width: int = 256, height: int = 256) -> bytes:
    """Create a valid JPEG image with given dimensions in memory."""
    img = Image.new("RGB", (width, height), color="blue")
    buf = io.BytesIO()
    img.save(buf, format="JPEG")
    return buf.getvalue()


def _create_printer(owner):
    """Create a test Printer record and return it."""
    from imptune.db.models import Printer

    return Printer.create(
        name="Test Printer",
        ip_address="10.0.0.1",
        port_name="IP_10.0.0.1",
        owner=owner,
    )


class TestIconUpload:
    def test_upload_valid_png(self, client, owner, tmp_data_dir):
        """POST a valid 256x256 PNG returns 200 and Icon record created."""
        from imptune.db.models import Icon

        printer = _create_printer(owner)
        png_data = _make_png(256, 256)
        response = client.post(
            f"/printers/{printer.id}/icon",
            files={"file": ("icon.png", io.BytesIO(png_data), "image/png")},
        )
        assert response.status_code == 200
        assert "Icon uploaded successfully" in response.text

        icons = list(Icon.select().where(Icon.printer == printer.id))
        assert len(icons) == 1
        assert icons[0].original_filename == "icon.png"

        # File must be stored on disk
        import hashlib
        from pathlib import Path

        sha256 = hashlib.sha256(png_data).hexdigest()
        icon_file = Path(tmp_data_dir) / "icons" / sha256
        assert icon_file.exists()

    def test_jpeg_is_converted(self, client, owner, tmp_data_dir):
        """A JPEG is accepted and re-encoded as a 256x256 PNG, not rejected."""
        from pathlib import Path

        from imptune.db.models import Icon

        printer = _create_printer(owner)
        jpeg_data = _make_jpeg(256, 256)
        response = client.post(
            f"/printers/{printer.id}/icon",
            files={"file": ("icon.jpg", io.BytesIO(jpeg_data), "image/jpeg")},
        )
        assert response.status_code == 200

        icon = Icon.get(Icon.printer == printer.id)
        stored = Path(tmp_data_dir) / "icons" / icon.sha256
        with Image.open(stored) as img:
            assert img.format == "PNG"
            assert img.size == (256, 256)

    def test_reject_undecodable_file(self, client, owner, tmp_data_dir):
        """A file Pillow cannot open is still refused."""
        printer = _create_printer(owner)
        response = client.post(
            f"/printers/{printer.id}/icon",
            files={"file": ("icon.png", io.BytesIO(b"not an image at all"), "image/png")},
        )
        assert response.status_code == 422
        assert "readable image" in response.text

    def test_reject_oversized(self, client, owner, tmp_data_dir):
        """POST with PNG > 750KB returns 422 with 750 KB error."""
        printer = _create_printer(owner)
        # Craft oversized data: valid PNG bytes followed by padding
        png_bytes = _make_png(256, 256)
        oversized = png_bytes + b"\x00" * (750 * 1024 + 1 - len(png_bytes))
        response = client.post(
            f"/printers/{printer.id}/icon",
            files={"file": ("big.png", io.BytesIO(oversized), "image/png")},
        )
        assert response.status_code == 422
        assert "750" in response.text

    def test_wrong_dimensions_are_resized(self, client, owner, tmp_data_dir):
        """An off-size PNG is letterboxed into 256x256 instead of rejected."""
        from pathlib import Path

        from imptune.db.models import Icon

        printer = _create_printer(owner)
        png_data = _make_png(128, 400)
        response = client.post(
            f"/printers/{printer.id}/icon",
            files={"file": ("tall.png", io.BytesIO(png_data), "image/png")},
        )
        assert response.status_code == 200

        icon = Icon.get(Icon.printer == printer.id)
        stored = Path(tmp_data_dir) / "icons" / icon.sha256
        with Image.open(stored) as img:
            assert img.size == (256, 256)

    def test_replace_existing_icon(self, client, owner, tmp_data_dir):
        """Second upload for same printer replaces the Icon record."""
        from imptune.db.models import Icon

        printer = _create_printer(owner)

        # First upload
        png1 = _make_png(256, 256)
        client.post(
            f"/printers/{printer.id}/icon",
            files={"file": ("icon1.png", io.BytesIO(png1), "image/png")},
        )

        # Second upload (different color to get different sha256)
        img2 = Image.new("RGBA", (256, 256), color="green")
        buf2 = io.BytesIO()
        img2.save(buf2, format="PNG")
        png2 = buf2.getvalue()

        response = client.post(
            f"/printers/{printer.id}/icon",
            files={"file": ("icon2.png", io.BytesIO(png2), "image/png")},
        )
        assert response.status_code == 200

        # Still only one Icon record for this printer
        icons = list(Icon.select().where(Icon.printer == printer.id))
        assert len(icons) == 1
        assert icons[0].original_filename == "icon2.png"

    def test_404_missing_printer(self, client, tmp_data_dir):
        """POST to nonexistent printer_id returns 404."""
        png_data = _make_png(256, 256)
        response = client.post(
            "/printers/99999/icon",
            files={"file": ("icon.png", io.BytesIO(png_data), "image/png")},
        )
        assert response.status_code == 404