96 lines
2.7 KiB
Python
96 lines
2.7 KiB
Python
"""Test initialization of the AI Task component."""
|
|
|
|
from freezegun.api import FrozenDateTimeFactory
|
|
import pytest
|
|
|
|
from homeassistant.components.ai_task import AITaskPreferences
|
|
from homeassistant.components.ai_task.const import DATA_PREFERENCES
|
|
from homeassistant.core import HomeAssistant
|
|
|
|
from .conftest import TEST_ENTITY_ID
|
|
|
|
from tests.common import flush_store
|
|
|
|
|
|
async def test_preferences_storage_load(
|
|
hass: HomeAssistant,
|
|
init_components: None,
|
|
freezer: FrozenDateTimeFactory,
|
|
) -> None:
|
|
"""Test that AITaskPreferences are stored and loaded correctly."""
|
|
preferences = hass.data[DATA_PREFERENCES]
|
|
|
|
# Initial state should be None for entity IDs
|
|
assert preferences.gen_text_entity_id is None
|
|
|
|
gen_text_id_1 = "sensor.summary_one"
|
|
|
|
preferences.async_set_preferences(
|
|
gen_text_entity_id=gen_text_id_1,
|
|
)
|
|
|
|
# Verify that current preferences object is updated
|
|
assert preferences.gen_text_entity_id == gen_text_id_1
|
|
|
|
await flush_store(preferences._store)
|
|
|
|
# Create a new preferences instance to test loading from store
|
|
new_preferences_instance = AITaskPreferences(hass)
|
|
await new_preferences_instance.async_load()
|
|
|
|
assert new_preferences_instance.gen_text_entity_id == gen_text_id_1
|
|
|
|
# Test updating one preference and setting another to None
|
|
gen_text_id_2 = "sensor.summary_two"
|
|
preferences.async_set_preferences(gen_text_entity_id=gen_text_id_2)
|
|
|
|
# Verify that current preferences object is updated
|
|
assert preferences.gen_text_entity_id == gen_text_id_2
|
|
|
|
await flush_store(preferences._store)
|
|
|
|
# Create another new preferences instance to confirm persistence of the update
|
|
another_new_preferences_instance = AITaskPreferences(hass)
|
|
await another_new_preferences_instance.async_load()
|
|
|
|
assert another_new_preferences_instance.gen_text_entity_id == gen_text_id_2
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("set_preferences", "msg_extra"),
|
|
[
|
|
(
|
|
{"gen_text_entity_id": TEST_ENTITY_ID},
|
|
{},
|
|
),
|
|
(
|
|
{},
|
|
{"entity_id": TEST_ENTITY_ID},
|
|
),
|
|
],
|
|
)
|
|
async def test_generate_text_service(
|
|
hass: HomeAssistant,
|
|
init_components: None,
|
|
freezer: FrozenDateTimeFactory,
|
|
set_preferences: dict[str, str | None],
|
|
msg_extra: dict[str, str],
|
|
) -> None:
|
|
"""Test the generate text service."""
|
|
preferences = hass.data[DATA_PREFERENCES]
|
|
preferences.async_set_preferences(**set_preferences)
|
|
|
|
result = await hass.services.async_call(
|
|
"ai_task",
|
|
"generate_text",
|
|
{
|
|
"task_name": "Test Name",
|
|
"instructions": "Test prompt",
|
|
}
|
|
| msg_extra,
|
|
blocking=True,
|
|
return_response=True,
|
|
)
|
|
|
|
assert result["result"] == "Mock result"
|