AI Task integration (#145128)

* Add AI Task integration * Remove GenTextTaskType * Add AI Task prefs * Add action to LLM task * Remove WS command * Rename result to text for GenTextTaskResult * Apply suggestions from code review Co-authored-by: Allen Porter <allen.porter@gmail.com> * Add supported feature for generate text * Update const.py Co-authored-by: HarvsG <11440490+HarvsG@users.noreply.github.com> * Update homeassistant/components/ai_task/services.yaml Co-authored-by: HarvsG <11440490+HarvsG@users.noreply.github.com> * Use WS API to set preferences * Simplify pref storage * Simplify pref test * Update homeassistant/components/ai_task/services.yaml Co-authored-by: Allen Porter <allen.porter@gmail.com> --------- Co-authored-by: Allen Porter <allen.porter@gmail.com> Co-authored-by: HarvsG <11440490+HarvsG@users.noreply.github.com>
2025-11-14 05:20:17 +00:00 · 2025-06-20 08:56:08 -04:00
parent 1b73acc025
commit 33bde48c9c
18 changed files with 922 additions and 0 deletions
--- a/tests/components/ai_task/test_task.py
+++ b/tests/components/ai_task/test_task.py
@@ -0,0 +1,123 @@
+"""Test tasks for the AI Task integration."""
+
+from freezegun import freeze_time
+import pytest
+from syrupy.assertion import SnapshotAssertion
+
+from homeassistant.components.ai_task import AITaskEntityFeature, async_generate_text
+from homeassistant.components.conversation import async_get_chat_log
+from homeassistant.const import STATE_UNKNOWN
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import chat_session
+
+from .conftest import TEST_ENTITY_ID, MockAITaskEntity
+
+from tests.typing import WebSocketGenerator
+
+
+async def test_run_task_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."""
+    client = await hass_ws_client(hass)
+
+    with pytest.raises(
+        ValueError, match="No entity_id provided and no preferred entity set"
+    ):
+        await async_generate_text(
+            hass,
+            task_name="Test Task",
+            instructions="Test prompt",
+        )
+
+    await client.send_json_auto_id(
+        {
+            "type": "ai_task/preferences/set",
+            "gen_text_entity_id": "ai_task.unknown",
+        }
+    )
+    msg = await client.receive_json()
+    assert msg["success"]
+
+    with pytest.raises(ValueError, match="AI Task entity ai_task.unknown not found"):
+        await async_generate_text(
+            hass,
+            task_name="Test Task",
+            instructions="Test prompt",
+        )
+
+    await client.send_json_auto_id(
+        {
+            "type": "ai_task/preferences/set",
+            "gen_text_entity_id": TEST_ENTITY_ID,
+        }
+    )
+    msg = await client.receive_json()
+    assert msg["success"]
+
+    state = hass.states.get(TEST_ENTITY_ID)
+    assert state is not None
+    assert state.state == STATE_UNKNOWN
+
+    result = await async_generate_text(
+        hass,
+        task_name="Test Task",
+        instructions="Test prompt",
+    )
+    assert result.text == "Mock result"
+    state = hass.states.get(TEST_ENTITY_ID)
+    assert state is not None
+    assert state.state != STATE_UNKNOWN
+
+    mock_ai_task_entity.supported_features = AITaskEntityFeature(0)
+    with pytest.raises(
+        ValueError,
+        match="AI Task entity ai_task.test_task_entity does not support generating text",
+    ):
+        await async_generate_text(
+            hass,
+            task_name="Test Task",
+            instructions="Test prompt",
+        )
+
+
+async def test_run_text_task_unknown_entity(
+    hass: HomeAssistant,
+    init_components: None,
+) -> None:
+    """Test running a text task with an unknown entity."""
+
+    with pytest.raises(
+        ValueError, match="AI Task entity ai_task.unknown_entity not found"
+    ):
+        await async_generate_text(
+            hass,
+            task_name="Test Task",
+            entity_id="ai_task.unknown_entity",
+            instructions="Test prompt",
+        )
+
+
+@freeze_time("2025-06-14 22:59:00")
+async def test_run_text_task_updates_chat_log(
+    hass: HomeAssistant,
+    init_components: None,
+    snapshot: SnapshotAssertion,
+) -> None:
+    """Test that running a text task updates the chat log."""
+    result = await async_generate_text(
+        hass,
+        task_name="Test Task",
+        entity_id=TEST_ENTITY_ID,
+        instructions="Test prompt",
+    )
+    assert result.text == "Mock result"
+
+    with (
+        chat_session.async_get_chat_session(hass, result.conversation_id) as session,
+        async_get_chat_log(hass, session) as chat_log,
+    ):
+        assert chat_log.content == snapshot