Add attachment support to AI task (#148120)

tests/components/ai_task/conftest.py

@@ -35,7 +35,9 @@ class MockAITaskEntity(AITaskEntity):
     """Mock AI Task entity for testing."""
 
     _attr_name = "Test Task Entity"
-    _attr_supported_features = AITaskEntityFeature.GENERATE_DATA
+    _attr_supported_features = (
+        AITaskEntityFeature.GENERATE_DATA | AITaskEntityFeature.SUPPORT_ATTACHMENTS
+    )
 
     def __init__(self) -> None:
         """Initialize the mock entity."""

tests/components/ai_task/test_init.py

@@ -1,11 +1,13 @@
 """Test initialization of the AI Task component."""
 
 from typing import Any
+from unittest.mock import patch
 
 from freezegun.api import FrozenDateTimeFactory
 import pytest
 import voluptuous as vol
 
+from homeassistant.components import media_source
 from homeassistant.components.ai_task import AITaskPreferences
 from homeassistant.components.ai_task.const import DATA_PREFERENCES
 from homeassistant.core import HomeAssistant
@@ -58,7 +60,15 @@ async def test_preferences_storage_load(
         ),
         (
             {},
-            {"entity_id": TEST_ENTITY_ID},
+            {
+                "entity_id": TEST_ENTITY_ID,
+                "attachments": [
+                    {
+                        "media_content_id": "media-source://mock/blah_blah_blah.mp4",
+                        "media_content_type": "video/mp4",
+                    }
+                ],
+            },
         ),
     ],
 )
@@ -68,25 +78,50 @@ async def test_generate_data_service(
     freezer: FrozenDateTimeFactory,
     set_preferences: dict[str, str | None],
     msg_extra: dict[str, str],
+    mock_ai_task_entity: MockAITaskEntity,
 ) -> None:
     """Test the generate data service."""
     preferences = hass.data[DATA_PREFERENCES]
     preferences.async_set_preferences(**set_preferences)
 
-    result = await hass.services.async_call(
-        "ai_task",
-        "generate_data",
-        {
-            "task_name": "Test Name",
-            "instructions": "Test prompt",
-        }
-        | msg_extra,
-        blocking=True,
-        return_response=True,
-    )
+    with patch(
+        "homeassistant.components.media_source.async_resolve_media",
+        return_value=media_source.PlayMedia(
+            url="http://example.com/media.mp4",
+            mime_type="video/mp4",
+        ),
+    ):
+        result = await hass.services.async_call(
+            "ai_task",
+            "generate_data",
+            {
+                "task_name": "Test Name",
+                "instructions": "Test prompt",
+            }
+            | msg_extra,
+            blocking=True,
+            return_response=True,
+        )
 
     assert result["data"] == "Mock result"
 
+    assert len(mock_ai_task_entity.mock_generate_data_tasks) == 1
+    task = mock_ai_task_entity.mock_generate_data_tasks[0]
+
+    assert len(task.attachments or []) == len(
+        msg_attachments := msg_extra.get("attachments", [])
+    )
+
+    for msg_attachment, attachment in zip(
+        msg_attachments, task.attachments or [], strict=False
+    ):
+        assert attachment.url == "http://example.com/media.mp4"
+        assert attachment.mime_type == "video/mp4"
+        assert attachment.media_content_id == msg_attachment["media_content_id"]
+        assert (
+            str(attachment) == f"<PlayMediaWithId {msg_attachment['media_content_id']}>"
+        )
+
 
 async def test_generate_data_service_structure_fields(
     hass: HomeAssistant,

tests/components/ai_task/test_task.py

@@ -16,13 +16,13 @@ from .conftest import TEST_ENTITY_ID, MockAITaskEntity
 from tests.typing import WebSocketGenerator
 
 
-async def test_run_task_preferred_entity(
+async def test_generate_data_preferred_entity(
     hass: HomeAssistant,
     init_components: None,
     mock_ai_task_entity: MockAITaskEntity,
     hass_ws_client: WebSocketGenerator,
 ) -> None:
-    """Test running a task with an unknown entity."""
+    """Test generating data with entity via preferences."""
     client = await hass_ws_client(hass)
 
     with pytest.raises(
@@ -90,11 +90,11 @@ async def test_run_task_preferred_entity(
         )
 
 
-async def test_run_data_task_unknown_entity(
+async def test_generate_data_unknown_entity(
     hass: HomeAssistant,
     init_components: None,
 ) -> None:
-    """Test running a data task with an unknown entity."""
+    """Test generating data with an unknown entity."""
 
     with pytest.raises(
         HomeAssistantError, match="AI Task entity ai_task.unknown_entity not found"
@@ -113,7 +113,7 @@ async def test_run_data_task_updates_chat_log(
     init_components: None,
     snapshot: SnapshotAssertion,
 ) -> None:
-    """Test that running a data task updates the chat log."""
+    """Test that generating data updates the chat log."""
     result = await async_generate_data(
         hass,
         task_name="Test Task",
@@ -127,3 +127,30 @@ async def test_run_data_task_updates_chat_log(
         async_get_chat_log(hass, session) as chat_log,
     ):
         assert chat_log.content == snapshot
+
+
+async def test_generate_data_attachments_not_supported(
+    hass: HomeAssistant,
+    init_components: None,
+    mock_ai_task_entity: MockAITaskEntity,
+) -> None:
+    """Test generating data with attachments when entity doesn't support them."""
+    # Remove attachment support from the entity
+    mock_ai_task_entity._attr_supported_features = AITaskEntityFeature.GENERATE_DATA
+
+    with pytest.raises(
+        HomeAssistantError,
+        match="AI Task entity ai_task.test_task_entity does not support attachments",
+    ):
+        await async_generate_data(
+            hass,
+            task_name="Test Task",
+            entity_id=TEST_ENTITY_ID,
+            instructions="Test prompt",
+            attachments=[
+                {
+                    "media_content_id": "media-source://mock/test.mp4",
+                    "media_content_type": "video/mp4",
+                }
+            ],
+        )