diff --git a/homeassistant/components/ai_task/task.py b/homeassistant/components/ai_task/task.py index e6d86bee978..1d27f75b6c7 100644 --- a/homeassistant/components/ai_task/task.py +++ b/homeassistant/components/ai_task/task.py @@ -12,7 +12,7 @@ from typing import Any import voluptuous as vol -from homeassistant.components import camera, conversation, media_source +from homeassistant.components import camera, conversation, image, media_source from homeassistant.components.http.auth import async_sign_path from homeassistant.core import HomeAssistant, ServiceResponse, callback from homeassistant.exceptions import HomeAssistantError @@ -31,14 +31,14 @@ from .const import ( ) -def _save_camera_snapshot(image: camera.Image) -> Path: +def _save_camera_snapshot(image_data: camera.Image | image.Image) -> Path: """Save camera snapshot to temp file.""" with tempfile.NamedTemporaryFile( mode="wb", - suffix=mimetypes.guess_extension(image.content_type, False), + suffix=mimetypes.guess_extension(image_data.content_type, False), delete=False, ) as temp_file: - temp_file.write(image.content) + temp_file.write(image_data.content) return Path(temp_file.name) @@ -54,26 +54,31 @@ async def _resolve_attachments( for attachment in attachments or []: media_content_id = attachment["media_content_id"] - # Special case for camera media sources - if media_content_id.startswith("media-source://camera/"): - # Extract entity_id from the media content ID - entity_id = media_content_id.removeprefix("media-source://camera/") + # Special case for certain media sources + for integration in camera, image: + media_source_prefix = f"media-source://{integration.DOMAIN}/" + if not media_content_id.startswith(media_source_prefix): + continue - # Get snapshot from camera - image = await camera.async_get_image(hass, entity_id) + # Extract entity_id from the media content ID + entity_id = media_content_id.removeprefix(media_source_prefix) + + # Get snapshot from entity + image_data = await integration.async_get_image(hass, entity_id) temp_filename = await hass.async_add_executor_job( - _save_camera_snapshot, image + _save_camera_snapshot, image_data ) created_files.append(temp_filename) resolved_attachments.append( conversation.Attachment( media_content_id=media_content_id, - mime_type=image.content_type, + mime_type=image_data.content_type, path=temp_filename, ) ) + break else: # Handle regular media sources media = await media_source.async_resolve_media(hass, media_content_id, None) diff --git a/homeassistant/components/image/__init__.py b/homeassistant/components/image/__init__.py index 0a3b9bf9af7..7bf0060f593 100644 --- a/homeassistant/components/image/__init__.py +++ b/homeassistant/components/image/__init__.py @@ -105,6 +105,20 @@ async def _async_get_image(image_entity: ImageEntity, timeout: int) -> Image: raise HomeAssistantError("Unable to get image") +async def async_get_image( + hass: HomeAssistant, + entity_id: str, + timeout: int = 10, +) -> Image: + """Fetch an image from an image entity.""" + component = hass.data[DATA_COMPONENT] + + if (image := component.get_entity(entity_id)) is None: + raise HomeAssistantError(f"Image entity {entity_id} not found") + + return await _async_get_image(image, timeout) + + async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: """Set up the image component.""" component = hass.data[DATA_COMPONENT] = EntityComponent[ImageEntity]( diff --git a/tests/components/ai_task/test_task.py b/tests/components/ai_task/test_task.py index 345d6c30981..4f8616d3f81 100644 --- a/tests/components/ai_task/test_task.py +++ b/tests/components/ai_task/test_task.py @@ -187,7 +187,11 @@ async def test_generate_data_mixed_attachments( patch( "homeassistant.components.camera.async_get_image", return_value=Image(content_type="image/jpeg", content=b"fake_camera_jpeg"), - ) as mock_get_image, + ) as mock_get_camera_image, + patch( + "homeassistant.components.image.async_get_image", + return_value=Image(content_type="image/jpeg", content=b"fake_image_jpeg"), + ) as mock_get_image_image, patch( "homeassistant.components.media_source.async_resolve_media", return_value=media_source.PlayMedia( @@ -207,6 +211,10 @@ async def test_generate_data_mixed_attachments( "media_content_id": "media-source://camera/camera.front_door", "media_content_type": "image/jpeg", }, + { + "media_content_id": "media-source://image/image.floorplan", + "media_content_type": "image/jpeg", + }, { "media_content_id": "media-source://media_player/video.mp4", "media_content_type": "video/mp4", @@ -215,7 +223,8 @@ async def test_generate_data_mixed_attachments( ) # Verify both methods were called - mock_get_image.assert_called_once_with(hass, "camera.front_door") + mock_get_camera_image.assert_called_once_with(hass, "camera.front_door") + mock_get_image_image.assert_called_once_with(hass, "image.floorplan") mock_resolve_media.assert_called_once_with( hass, "media-source://media_player/video.mp4", None ) @@ -224,7 +233,7 @@ async def test_generate_data_mixed_attachments( assert len(mock_ai_task_entity.mock_generate_data_tasks) == 1 task = mock_ai_task_entity.mock_generate_data_tasks[0] assert task.attachments is not None - assert len(task.attachments) == 2 + assert len(task.attachments) == 3 # Check camera attachment camera_attachment = task.attachments[0] @@ -240,6 +249,18 @@ async def test_generate_data_mixed_attachments( content = await hass.async_add_executor_job(camera_attachment.path.read_bytes) assert content == b"fake_camera_jpeg" + # Check image attachment + image_attachment = task.attachments[1] + assert image_attachment.media_content_id == "media-source://image/image.floorplan" + assert image_attachment.mime_type == "image/jpeg" + assert isinstance(image_attachment.path, Path) + assert image_attachment.path.suffix == ".jpg" + + # Verify image snapshot content + assert image_attachment.path.exists() + content = await hass.async_add_executor_job(image_attachment.path.read_bytes) + assert content == b"fake_image_jpeg" + # Trigger clean up async_fire_time_changed( hass, @@ -249,9 +270,10 @@ async def test_generate_data_mixed_attachments( # Verify the temporary file cleaned up assert not camera_attachment.path.exists() + assert not image_attachment.path.exists() # Check regular media attachment - media_attachment = task.attachments[1] + media_attachment = task.attachments[2] assert media_attachment.media_content_id == "media-source://media_player/video.mp4" assert media_attachment.mime_type == "video/mp4" assert media_attachment.path == Path("/media/test.mp4") diff --git a/tests/components/image/test_init.py b/tests/components/image/test_init.py index bb8762f17e2..0a1c939c474 100644 --- a/tests/components/image/test_init.py +++ b/tests/components/image/test_init.py @@ -407,6 +407,15 @@ async def test_image_stream( await close_future +async def test_get_image_action(hass: HomeAssistant, mock_image_platform: None) -> None: + """Test get_image action.""" + image_data = await image.async_get_image(hass, "image.test") + assert image_data == image.Image(content_type="image/jpeg", content=b"Test") + + with pytest.raises(HomeAssistantError, match="not found"): + await image.async_get_image(hass, "image.unknown") + + async def test_snapshot_service(hass: HomeAssistant) -> None: """Test snapshot service.""" mopen = mock_open()