"""Test that we can upload images.""" import pathlib import tempfile from unittest.mock import patch from aiohttp import ClientSession, ClientWebSocketResponse from freezegun.api import FrozenDateTimeFactory from PIL import Image import pytest from homeassistant.components.image_upload import DOMAIN from homeassistant.components.websocket_api import TYPE_RESULT from homeassistant.core import HomeAssistant from homeassistant.setup import async_setup_component from homeassistant.util import dt as dt_util from . import TEST_IMAGE from tests.typing import ClientSessionGenerator, WebSocketGenerator async def test_upload_image( hass: HomeAssistant, freezer: FrozenDateTimeFactory, hass_client: ClientSessionGenerator, hass_ws_client: WebSocketGenerator, ) -> None: """Test we can upload an image.""" now = dt_util.utcnow() freezer.move_to(now) with ( tempfile.TemporaryDirectory() as tempdir, patch.object(hass.config, "path", return_value=tempdir), ): assert await async_setup_component(hass, DOMAIN, {}) ws_client: ClientWebSocketResponse = await hass_ws_client() client: ClientSession = await hass_client() with TEST_IMAGE.open("rb") as fp: res = await client.post("/api/image/upload", data={"file": fp}) assert res.status == 200 item = await res.json() assert item["content_type"] == "image/png" assert item["filesize"] == 38847 assert item["name"] == "logo.png" assert item["uploaded_at"] == now.isoformat() tempdir = pathlib.Path(tempdir) item_folder: pathlib.Path = tempdir / item["id"] test_image_bytes = TEST_IMAGE.read_bytes() assert (item_folder / "original").read_bytes() == test_image_bytes # fetch original image res = await client.get(f"/api/image/serve/{item['id']}/original") assert res.status == 200 fetched_image_bytes = await res.read() assert fetched_image_bytes == test_image_bytes # fetch non-existing image res = await client.get("/api/image/serve/non-existing/256x256") assert res.status == 404 # fetch invalid sizes for inv_size in ("256", "256x25A", "100x100", "25Ax256"): res = await client.get(f"/api/image/serve/{item['id']}/{inv_size}") assert res.status == 400 # fetch resized version res = await client.get(f"/api/image/serve/{item['id']}/256x256") assert res.status == 200 assert (item_folder / "256x256").is_file() # List item await ws_client.send_json({"id": 6, "type": "image/list"}) msg = await ws_client.receive_json() assert msg["id"] == 6 assert msg["type"] == TYPE_RESULT assert msg["success"] assert msg["result"] == [item] # Delete item await ws_client.send_json( {"id": 7, "type": "image/delete", "image_id": item["id"]} ) msg = await ws_client.receive_json() assert msg["id"] == 7 assert msg["type"] == TYPE_RESULT assert msg["success"] # Ensure removed from disk assert not item_folder.is_dir() @pytest.mark.parametrize( ("image_mode", "content_type"), [ ("RGBA", "image/jpeg"), ("LA", "image/jpeg"), ("P", "image/jpeg"), ], ) async def test_upload_image_thumbnail_rgba_as_jpeg( hass: HomeAssistant, freezer: FrozenDateTimeFactory, hass_client: ClientSessionGenerator, image_mode: str, content_type: str, ) -> None: """Test thumbnail generation when image mode is incompatible with JPEG.""" now = dt_util.utcnow() freezer.move_to(now) with ( tempfile.TemporaryDirectory() as tempdir, patch.object(hass.config, "path", return_value=tempdir), ): assert await async_setup_component(hass, DOMAIN, {}) client: ClientSession = await hass_client() with TEST_IMAGE.open("rb") as fp: res = await client.post("/api/image/upload", data={"file": fp}) assert res.status == 200 item = await res.json() image_id = item["id"] tempdir = pathlib.Path(tempdir) item_folder = tempdir / image_id # Create an image file with the given mode to simulate the mismatch original_path = item_folder / "original" img = Image.new(image_mode, (300, 300)) img.save(original_path, format="png") # Change the stored content_type to simulate the mismatch hass.data[DOMAIN].data[image_id]["content_type"] = content_type # Fetch the thumbnail; this should not raise an OSError res = await client.get(f"/api/image/serve/{image_id}/256x256") assert res.status == 200 assert (item_folder / "256x256").is_file() # Verify the generated thumbnail is a valid JPEG thumbnail = Image.open(item_folder / "256x256") assert thumbnail.mode == "RGB"