"""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