79 lines
2.5 KiB
Python
79 lines
2.5 KiB
Python
"""Tests for the AI Task entity model."""
|
|
|
|
from freezegun import freeze_time
|
|
import voluptuous as vol
|
|
|
|
from homeassistant.components.ai_task import async_generate_data
|
|
from homeassistant.const import STATE_UNKNOWN
|
|
from homeassistant.core import HomeAssistant
|
|
from homeassistant.helpers import selector
|
|
|
|
from .conftest import TEST_ENTITY_ID, MockAITaskEntity
|
|
|
|
from tests.common import MockConfigEntry
|
|
|
|
|
|
@freeze_time("2025-06-08 16:28:13")
|
|
async def test_state_generate_data(
|
|
hass: HomeAssistant,
|
|
init_components: None,
|
|
mock_config_entry: MockConfigEntry,
|
|
mock_ai_task_entity: MockAITaskEntity,
|
|
) -> None:
|
|
"""Test the state of the AI Task entity is updated when generating data."""
|
|
entity = hass.states.get(TEST_ENTITY_ID)
|
|
assert entity is not None
|
|
assert entity.state == STATE_UNKNOWN
|
|
|
|
result = await async_generate_data(
|
|
hass,
|
|
task_name="Test task",
|
|
entity_id=TEST_ENTITY_ID,
|
|
instructions="Test prompt",
|
|
)
|
|
assert result.data == "Mock result"
|
|
|
|
entity = hass.states.get(TEST_ENTITY_ID)
|
|
assert entity.state == "2025-06-08T16:28:13+00:00"
|
|
|
|
assert mock_ai_task_entity.mock_generate_data_tasks
|
|
task = mock_ai_task_entity.mock_generate_data_tasks[0]
|
|
assert task.instructions == "Test prompt"
|
|
|
|
|
|
async def test_generate_structured_data(
|
|
hass: HomeAssistant,
|
|
init_components: None,
|
|
mock_config_entry: MockConfigEntry,
|
|
mock_ai_task_entity: MockAITaskEntity,
|
|
) -> None:
|
|
"""Test the entity can generate structured data."""
|
|
result = await async_generate_data(
|
|
hass,
|
|
task_name="Test task",
|
|
entity_id=TEST_ENTITY_ID,
|
|
instructions="Please generate a profile for a new user",
|
|
structure=vol.Schema(
|
|
{
|
|
vol.Required("name"): selector.TextSelector(),
|
|
vol.Optional("age"): selector.NumberSelector(
|
|
config=selector.NumberSelectorConfig(
|
|
min=0,
|
|
max=120,
|
|
)
|
|
),
|
|
}
|
|
),
|
|
)
|
|
# Arbitrary data returned by the mock entity (not determined by above schema in test)
|
|
assert result.data == {
|
|
"name": "Tracy Chen",
|
|
"age": 30,
|
|
}
|
|
|
|
assert mock_ai_task_entity.mock_generate_data_tasks
|
|
task = mock_ai_task_entity.mock_generate_data_tasks[0]
|
|
assert task.instructions == "Please generate a profile for a new user"
|
|
assert task.structure
|
|
assert isinstance(task.structure, vol.Schema)
|