diff --git a/homeassistant/components/image/__init__.py b/homeassistant/components/image/__init__.py index 08f8431556c..72636e3bf4a 100644 --- a/homeassistant/components/image/__init__.py +++ b/homeassistant/components/image/__init__.py @@ -58,12 +58,23 @@ class Image: content: bytes +class ImageContentTypeError(HomeAssistantError): + """Error with the content type while loading an image.""" + + +def valid_image_content_type(content_type: str) -> str: + """Validate the assigned content type is one of an image.""" + if content_type.split("/", 1)[0] != "image": + raise ImageContentTypeError + return content_type + + async def _async_get_image(image_entity: ImageEntity, timeout: int) -> Image: """Fetch image from an image entity.""" - with suppress(asyncio.CancelledError, asyncio.TimeoutError): + with suppress(asyncio.CancelledError, asyncio.TimeoutError, ImageContentTypeError): async with async_timeout.timeout(timeout): if image_bytes := await image_entity.async_image(): - content_type = image_entity.content_type + content_type = valid_image_content_type(image_entity.content_type) image = Image(content_type, image_bytes) return image diff --git a/tests/components/image/conftest.py b/tests/components/image/conftest.py index 38d386eeb6b..03ef6d03f59 100644 --- a/tests/components/image/conftest.py +++ b/tests/components/image/conftest.py @@ -36,6 +36,21 @@ class MockImageEntity(image.ImageEntity): return b"Test" +class MockImageEntityInvalidContentType(image.ImageEntity): + """Mock image entity.""" + + _attr_name = "Test" + + async def async_added_to_hass(self): + """Set the update time and assign and incorrect content type.""" + self._attr_content_type = "text/json" + self._attr_image_last_updated = dt_util.utcnow() + + async def async_image(self) -> bytes | None: + """Return bytes of image.""" + return b"Test" + + class MockURLImageEntity(image.ImageEntity): """Mock image entity.""" @@ -125,7 +140,9 @@ def config_flow_fixture(hass: HomeAssistant) -> Generator[None, None, None]: @pytest.fixture(name="mock_image_config_entry") -async def mock_image_config_entry_fixture(hass: HomeAssistant, config_flow: None): +async def mock_image_config_entry_fixture( + hass: HomeAssistant, config_flow: None +) -> ConfigEntry: """Initialize a mock image config_entry.""" async def async_setup_entry_init( @@ -166,7 +183,7 @@ async def mock_image_config_entry_fixture(hass: HomeAssistant, config_flow: None @pytest.fixture(name="mock_image_platform") -async def mock_image_platform_fixture(hass: HomeAssistant): +async def mock_image_platform_fixture(hass: HomeAssistant) -> None: """Initialize a mock image platform.""" mock_integration(hass, MockModule(domain="test")) mock_platform(hass, "test.image", MockImagePlatform([MockImageEntity(hass)])) diff --git a/tests/components/image/test_init.py b/tests/components/image/test_init.py index 3555026d45c..0573484a534 100644 --- a/tests/components/image/test_init.py +++ b/tests/components/image/test_init.py @@ -9,11 +9,13 @@ import pytest import respx from homeassistant.components import image +from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.setup import async_setup_component from .conftest import ( MockImageEntity, + MockImageEntityInvalidContentType, MockImageNoStateEntity, MockImagePlatform, MockImageSyncEntity, @@ -26,7 +28,7 @@ from tests.typing import ClientSessionGenerator @pytest.mark.freeze_time("2023-04-01 00:00:00+00:00") async def test_state( - hass: HomeAssistant, hass_client: ClientSessionGenerator, mock_image_platform + hass: HomeAssistant, hass_client: ClientSessionGenerator, mock_image_platform: None ) -> None: """Test image state.""" state = hass.states.get("image.test") @@ -41,7 +43,9 @@ async def test_state( @pytest.mark.freeze_time("2023-04-01 00:00:00+00:00") async def test_config_entry( - hass: HomeAssistant, hass_client: ClientSessionGenerator, mock_image_config_entry + hass: HomeAssistant, + hass_client: ClientSessionGenerator, + mock_image_config_entry: ConfigEntry, ) -> None: """Test setting up an image platform from a config entry.""" state = hass.states.get("image.test") @@ -99,8 +103,35 @@ async def test_no_state( } +async def test_no_valid_content_type( + hass: HomeAssistant, hass_client: ClientSessionGenerator +) -> None: + """Test invalid content type.""" + mock_integration(hass, MockModule(domain="test")) + mock_platform( + hass, "test.image", MockImagePlatform([MockImageEntityInvalidContentType(hass)]) + ) + assert await async_setup_component( + hass, image.DOMAIN, {"image": {"platform": "test"}} + ) + await hass.async_block_till_done() + + client = await hass_client() + + state = hass.states.get("image.test") + # assert state.state == "unknown" + access_token = state.attributes["access_token"] + assert state.attributes == { + "access_token": access_token, + "entity_picture": f"/api/image_proxy/image.test?token={access_token}", + "friendly_name": "Test", + } + resp = await client.get(f"/api/image_proxy/image.test?token={access_token}") + assert resp.status == HTTPStatus.INTERNAL_SERVER_ERROR + + async def test_fetch_image_authenticated( - hass: HomeAssistant, hass_client: ClientSessionGenerator, mock_image_platform + hass: HomeAssistant, hass_client: ClientSessionGenerator, mock_image_platform: None ) -> None: """Test fetching an image with an authenticated client.""" client = await hass_client() @@ -115,7 +146,7 @@ async def test_fetch_image_authenticated( async def test_fetch_image_fail( - hass: HomeAssistant, hass_client: ClientSessionGenerator, mock_image_platform + hass: HomeAssistant, hass_client: ClientSessionGenerator, mock_image_platform: None ) -> None: """Test fetching an image with an authenticated client.""" client = await hass_client() @@ -147,7 +178,7 @@ async def test_fetch_image_sync( async def test_fetch_image_unauthenticated( hass: HomeAssistant, hass_client_no_auth: ClientSessionGenerator, - mock_image_platform, + mock_image_platform: None, ) -> None: """Test fetching an image with an unauthenticated client.""" client = await hass_client_no_auth()