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_reject_non_png(self, client, owner, tmp_data_dir):
        """POST with a JPEG file returns 422 with PNG format error."""
        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 == 422
        assert "PNG" 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_reject_wrong_dimensions(self, client, owner, tmp_data_dir):
        """POST with 128x128 PNG returns 422 with 256x256 error."""
        printer = _create_printer(owner)
        png_data = _make_png(128, 128)
        response = client.post(
            f"/printers/{printer.id}/icon",
            files={"file": ("small.png", io.BytesIO(png_data), "image/png")},
        )
        assert response.status_code == 422
        assert "256x256" in response.text

    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