mirror of
https://github.com/home-assistant/core.git
synced 2025-08-01 03:35:09 +02:00
Convert Ollama to subentries (#147286)
* Convert Ollama to subentries * Add latest changes from Google subentries * Move config entry type to init
This commit is contained in:
@@ -30,7 +30,15 @@ def mock_config_entry(
|
||||
entry = MockConfigEntry(
|
||||
domain=ollama.DOMAIN,
|
||||
data=TEST_USER_DATA,
|
||||
options=mock_config_entry_options,
|
||||
version=2,
|
||||
subentries_data=[
|
||||
{
|
||||
"data": mock_config_entry_options,
|
||||
"subentry_type": "conversation",
|
||||
"title": "Ollama Conversation",
|
||||
"unique_id": None,
|
||||
}
|
||||
],
|
||||
)
|
||||
entry.add_to_hass(hass)
|
||||
return entry
|
||||
@@ -41,8 +49,10 @@ def mock_config_entry_with_assist(
|
||||
hass: HomeAssistant, mock_config_entry: MockConfigEntry
|
||||
) -> MockConfigEntry:
|
||||
"""Mock a config entry with assist."""
|
||||
hass.config_entries.async_update_entry(
|
||||
mock_config_entry, options={CONF_LLM_HASS_API: llm.LLM_API_ASSIST}
|
||||
hass.config_entries.async_update_subentry(
|
||||
mock_config_entry,
|
||||
next(iter(mock_config_entry.subentries.values())),
|
||||
data={CONF_LLM_HASS_API: llm.LLM_API_ASSIST},
|
||||
)
|
||||
return mock_config_entry
|
||||
|
||||
|
@@ -63,6 +63,37 @@ async def test_form(hass: HomeAssistant) -> None:
|
||||
assert len(mock_setup_entry.mock_calls) == 1
|
||||
|
||||
|
||||
async def test_duplicate_entry(hass: HomeAssistant) -> None:
|
||||
"""Test we abort on duplicate config entry."""
|
||||
MockConfigEntry(
|
||||
domain=ollama.DOMAIN,
|
||||
data={
|
||||
ollama.CONF_URL: "http://localhost:11434",
|
||||
ollama.CONF_MODEL: "test_model",
|
||||
},
|
||||
).add_to_hass(hass)
|
||||
|
||||
result = await hass.config_entries.flow.async_init(
|
||||
ollama.DOMAIN, context={"source": config_entries.SOURCE_USER}
|
||||
)
|
||||
assert result["type"] is FlowResultType.FORM
|
||||
assert not result["errors"]
|
||||
|
||||
with patch(
|
||||
"homeassistant.components.ollama.config_flow.ollama.AsyncClient.list",
|
||||
return_value={"models": [{"model": "test_model"}]},
|
||||
):
|
||||
result = await hass.config_entries.flow.async_configure(
|
||||
result["flow_id"],
|
||||
{
|
||||
ollama.CONF_URL: "http://localhost:11434",
|
||||
},
|
||||
)
|
||||
|
||||
assert result["type"] is FlowResultType.ABORT
|
||||
assert result["reason"] == "already_configured"
|
||||
|
||||
|
||||
async def test_form_need_download(hass: HomeAssistant) -> None:
|
||||
"""Test flow when a model needs to be downloaded."""
|
||||
# Pretend we already set up a config entry.
|
||||
@@ -155,14 +186,21 @@ async def test_form_need_download(hass: HomeAssistant) -> None:
|
||||
assert len(mock_setup_entry.mock_calls) == 1
|
||||
|
||||
|
||||
async def test_options(
|
||||
async def test_subentry_options(
|
||||
hass: HomeAssistant, mock_config_entry, mock_init_component
|
||||
) -> None:
|
||||
"""Test the options form."""
|
||||
options_flow = await hass.config_entries.options.async_init(
|
||||
mock_config_entry.entry_id
|
||||
"""Test the subentry options form."""
|
||||
subentry = next(iter(mock_config_entry.subentries.values()))
|
||||
|
||||
# Test reconfiguration
|
||||
options_flow = await mock_config_entry.start_subentry_reconfigure_flow(
|
||||
hass, subentry.subentry_id
|
||||
)
|
||||
options = await hass.config_entries.options.async_configure(
|
||||
|
||||
assert options_flow["type"] is FlowResultType.FORM
|
||||
assert options_flow["step_id"] == "set_options"
|
||||
|
||||
options = await hass.config_entries.subentries.async_configure(
|
||||
options_flow["flow_id"],
|
||||
{
|
||||
ollama.CONF_PROMPT: "test prompt",
|
||||
@@ -172,8 +210,10 @@ async def test_options(
|
||||
},
|
||||
)
|
||||
await hass.async_block_till_done()
|
||||
assert options["type"] is FlowResultType.CREATE_ENTRY
|
||||
assert options["data"] == {
|
||||
|
||||
assert options["type"] is FlowResultType.ABORT
|
||||
assert options["reason"] == "reconfigure_successful"
|
||||
assert subentry.data == {
|
||||
ollama.CONF_PROMPT: "test prompt",
|
||||
ollama.CONF_MAX_HISTORY: 100,
|
||||
ollama.CONF_NUM_CTX: 32768,
|
||||
@@ -181,6 +221,22 @@ async def test_options(
|
||||
}
|
||||
|
||||
|
||||
async def test_creating_conversation_subentry_not_loaded(
|
||||
hass: HomeAssistant,
|
||||
mock_init_component,
|
||||
mock_config_entry: MockConfigEntry,
|
||||
) -> None:
|
||||
"""Test creating a conversation subentry when entry is not loaded."""
|
||||
await hass.config_entries.async_unload(mock_config_entry.entry_id)
|
||||
result = await hass.config_entries.subentries.async_init(
|
||||
(mock_config_entry.entry_id, "conversation"),
|
||||
context={"source": config_entries.SOURCE_USER},
|
||||
)
|
||||
|
||||
assert result["type"] is FlowResultType.ABORT
|
||||
assert result["reason"] == "entry_not_loaded"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("side_effect", "error"),
|
||||
[
|
||||
|
@@ -35,7 +35,7 @@ async def stream_generator(response: dict | list[dict]) -> AsyncGenerator[dict]:
|
||||
yield msg
|
||||
|
||||
|
||||
@pytest.mark.parametrize("agent_id", [None, "conversation.mock_title"])
|
||||
@pytest.mark.parametrize("agent_id", [None, "conversation.ollama_conversation"])
|
||||
async def test_chat(
|
||||
hass: HomeAssistant,
|
||||
mock_config_entry: MockConfigEntry,
|
||||
@@ -149,9 +149,11 @@ async def test_template_variables(
|
||||
mock_user.id = "12345"
|
||||
mock_user.name = "Test User"
|
||||
|
||||
hass.config_entries.async_update_entry(
|
||||
subentry = next(iter(mock_config_entry.subentries.values()))
|
||||
hass.config_entries.async_update_subentry(
|
||||
mock_config_entry,
|
||||
options={
|
||||
subentry,
|
||||
data={
|
||||
"prompt": (
|
||||
"The user name is {{ user_name }}. "
|
||||
"The user id is {{ llm_context.context.user_id }}."
|
||||
@@ -382,10 +384,12 @@ async def test_unknown_hass_api(
|
||||
mock_init_component,
|
||||
) -> None:
|
||||
"""Test when we reference an API that no longer exists."""
|
||||
hass.config_entries.async_update_entry(
|
||||
subentry = next(iter(mock_config_entry.subentries.values()))
|
||||
hass.config_entries.async_update_subentry(
|
||||
mock_config_entry,
|
||||
options={
|
||||
**mock_config_entry.options,
|
||||
subentry,
|
||||
data={
|
||||
**subentry.data,
|
||||
CONF_LLM_HASS_API: "non-existing",
|
||||
},
|
||||
)
|
||||
@@ -518,8 +522,9 @@ async def test_message_history_unlimited(
|
||||
with (
|
||||
patch("ollama.AsyncClient.chat", side_effect=stream) as mock_chat,
|
||||
):
|
||||
hass.config_entries.async_update_entry(
|
||||
mock_config_entry, options={ollama.CONF_MAX_HISTORY: 0}
|
||||
subentry = next(iter(mock_config_entry.subentries.values()))
|
||||
hass.config_entries.async_update_subentry(
|
||||
mock_config_entry, subentry, data={ollama.CONF_MAX_HISTORY: 0}
|
||||
)
|
||||
for i in range(100):
|
||||
result = await conversation.async_converse(
|
||||
@@ -563,9 +568,11 @@ async def test_template_error(
|
||||
hass: HomeAssistant, mock_config_entry: MockConfigEntry
|
||||
) -> None:
|
||||
"""Test that template error handling works."""
|
||||
hass.config_entries.async_update_entry(
|
||||
subentry = next(iter(mock_config_entry.subentries.values()))
|
||||
hass.config_entries.async_update_subentry(
|
||||
mock_config_entry,
|
||||
options={
|
||||
subentry,
|
||||
data={
|
||||
"prompt": "talk like a {% if True %}smarthome{% else %}pirate please.",
|
||||
},
|
||||
)
|
||||
@@ -593,7 +600,7 @@ async def test_conversation_agent(
|
||||
)
|
||||
assert agent.supported_languages == MATCH_ALL
|
||||
|
||||
state = hass.states.get("conversation.mock_title")
|
||||
state = hass.states.get("conversation.ollama_conversation")
|
||||
assert state
|
||||
assert state.attributes[ATTR_SUPPORTED_FEATURES] == 0
|
||||
|
||||
@@ -609,7 +616,7 @@ async def test_conversation_agent_with_assist(
|
||||
)
|
||||
assert agent.supported_languages == MATCH_ALL
|
||||
|
||||
state = hass.states.get("conversation.mock_title")
|
||||
state = hass.states.get("conversation.ollama_conversation")
|
||||
assert state
|
||||
assert (
|
||||
state.attributes[ATTR_SUPPORTED_FEATURES]
|
||||
@@ -642,7 +649,7 @@ async def test_options(
|
||||
"test message",
|
||||
None,
|
||||
Context(),
|
||||
agent_id="conversation.mock_title",
|
||||
agent_id="conversation.ollama_conversation",
|
||||
)
|
||||
|
||||
assert mock_chat.call_count == 1
|
||||
@@ -667,9 +674,11 @@ async def test_reasoning_filter(
|
||||
entry = MockConfigEntry()
|
||||
entry.add_to_hass(hass)
|
||||
|
||||
hass.config_entries.async_update_entry(
|
||||
subentry = next(iter(mock_config_entry.subentries.values()))
|
||||
hass.config_entries.async_update_subentry(
|
||||
mock_config_entry,
|
||||
options={
|
||||
subentry,
|
||||
data={
|
||||
ollama.CONF_THINK: think,
|
||||
},
|
||||
)
|
||||
|
@@ -6,9 +6,13 @@ from httpx import ConnectError
|
||||
import pytest
|
||||
|
||||
from homeassistant.components import ollama
|
||||
from homeassistant.components.ollama.const import DOMAIN
|
||||
from homeassistant.core import HomeAssistant
|
||||
from homeassistant.helpers import device_registry as dr, entity_registry as er
|
||||
from homeassistant.setup import async_setup_component
|
||||
|
||||
from . import TEST_OPTIONS, TEST_USER_DATA
|
||||
|
||||
from tests.common import MockConfigEntry
|
||||
|
||||
|
||||
@@ -34,3 +38,250 @@ async def test_init_error(
|
||||
assert await async_setup_component(hass, ollama.DOMAIN, {})
|
||||
await hass.async_block_till_done()
|
||||
assert error in caplog.text
|
||||
|
||||
|
||||
async def test_migration_from_v1_to_v2(
|
||||
hass: HomeAssistant,
|
||||
device_registry: dr.DeviceRegistry,
|
||||
entity_registry: er.EntityRegistry,
|
||||
) -> None:
|
||||
"""Test migration from version 1 to version 2."""
|
||||
# Create a v1 config entry with conversation options and an entity
|
||||
mock_config_entry = MockConfigEntry(
|
||||
domain=DOMAIN,
|
||||
data=TEST_USER_DATA,
|
||||
options=TEST_OPTIONS,
|
||||
version=1,
|
||||
title="llama-3.2-8b",
|
||||
)
|
||||
mock_config_entry.add_to_hass(hass)
|
||||
|
||||
device = device_registry.async_get_or_create(
|
||||
config_entry_id=mock_config_entry.entry_id,
|
||||
identifiers={(DOMAIN, mock_config_entry.entry_id)},
|
||||
name=mock_config_entry.title,
|
||||
manufacturer="Ollama",
|
||||
model="Ollama",
|
||||
entry_type=dr.DeviceEntryType.SERVICE,
|
||||
)
|
||||
entity = entity_registry.async_get_or_create(
|
||||
"conversation",
|
||||
DOMAIN,
|
||||
mock_config_entry.entry_id,
|
||||
config_entry=mock_config_entry,
|
||||
device_id=device.id,
|
||||
suggested_object_id="llama_3_2_8b",
|
||||
)
|
||||
|
||||
# Run migration
|
||||
with patch(
|
||||
"homeassistant.components.ollama.async_setup_entry",
|
||||
return_value=True,
|
||||
):
|
||||
await hass.config_entries.async_setup(mock_config_entry.entry_id)
|
||||
|
||||
assert mock_config_entry.version == 2
|
||||
assert mock_config_entry.data == TEST_USER_DATA
|
||||
assert mock_config_entry.options == {}
|
||||
|
||||
assert len(mock_config_entry.subentries) == 1
|
||||
|
||||
subentry = next(iter(mock_config_entry.subentries.values()))
|
||||
assert subentry.unique_id is None
|
||||
assert subentry.title == "llama-3.2-8b"
|
||||
assert subentry.subentry_type == "conversation"
|
||||
assert subentry.data == TEST_OPTIONS
|
||||
|
||||
migrated_entity = entity_registry.async_get(entity.entity_id)
|
||||
assert migrated_entity is not None
|
||||
assert migrated_entity.config_entry_id == mock_config_entry.entry_id
|
||||
assert migrated_entity.config_subentry_id == subentry.subentry_id
|
||||
assert migrated_entity.unique_id == subentry.subentry_id
|
||||
|
||||
# Check device migration
|
||||
assert not device_registry.async_get_device(
|
||||
identifiers={(DOMAIN, mock_config_entry.entry_id)}
|
||||
)
|
||||
assert (
|
||||
migrated_device := device_registry.async_get_device(
|
||||
identifiers={(DOMAIN, subentry.subentry_id)}
|
||||
)
|
||||
)
|
||||
assert migrated_device.identifiers == {(DOMAIN, subentry.subentry_id)}
|
||||
assert migrated_device.id == device.id
|
||||
|
||||
|
||||
async def test_migration_from_v1_to_v2_with_multiple_urls(
|
||||
hass: HomeAssistant,
|
||||
device_registry: dr.DeviceRegistry,
|
||||
entity_registry: er.EntityRegistry,
|
||||
) -> None:
|
||||
"""Test migration from version 1 to version 2 with different URLs."""
|
||||
# Create two v1 config entries with different URLs
|
||||
mock_config_entry = MockConfigEntry(
|
||||
domain=DOMAIN,
|
||||
data={"url": "http://localhost:11434", "model": "llama3.2:latest"},
|
||||
options=TEST_OPTIONS,
|
||||
version=1,
|
||||
title="Ollama 1",
|
||||
)
|
||||
mock_config_entry.add_to_hass(hass)
|
||||
mock_config_entry_2 = MockConfigEntry(
|
||||
domain=DOMAIN,
|
||||
data={"url": "http://localhost:11435", "model": "llama3.2:latest"},
|
||||
options=TEST_OPTIONS,
|
||||
version=1,
|
||||
title="Ollama 2",
|
||||
)
|
||||
mock_config_entry_2.add_to_hass(hass)
|
||||
|
||||
device = device_registry.async_get_or_create(
|
||||
config_entry_id=mock_config_entry.entry_id,
|
||||
identifiers={(DOMAIN, mock_config_entry.entry_id)},
|
||||
name=mock_config_entry.title,
|
||||
manufacturer="Ollama",
|
||||
model="Ollama 1",
|
||||
entry_type=dr.DeviceEntryType.SERVICE,
|
||||
)
|
||||
entity_registry.async_get_or_create(
|
||||
"conversation",
|
||||
DOMAIN,
|
||||
mock_config_entry.entry_id,
|
||||
config_entry=mock_config_entry,
|
||||
device_id=device.id,
|
||||
suggested_object_id="ollama_1",
|
||||
)
|
||||
|
||||
device_2 = device_registry.async_get_or_create(
|
||||
config_entry_id=mock_config_entry_2.entry_id,
|
||||
identifiers={(DOMAIN, mock_config_entry_2.entry_id)},
|
||||
name=mock_config_entry_2.title,
|
||||
manufacturer="Ollama",
|
||||
model="Ollama 2",
|
||||
entry_type=dr.DeviceEntryType.SERVICE,
|
||||
)
|
||||
entity_registry.async_get_or_create(
|
||||
"conversation",
|
||||
DOMAIN,
|
||||
mock_config_entry_2.entry_id,
|
||||
config_entry=mock_config_entry_2,
|
||||
device_id=device_2.id,
|
||||
suggested_object_id="ollama_2",
|
||||
)
|
||||
|
||||
# Run migration
|
||||
with patch(
|
||||
"homeassistant.components.ollama.async_setup_entry",
|
||||
return_value=True,
|
||||
):
|
||||
await hass.config_entries.async_setup(mock_config_entry.entry_id)
|
||||
await hass.async_block_till_done()
|
||||
|
||||
entries = hass.config_entries.async_entries(DOMAIN)
|
||||
assert len(entries) == 2
|
||||
|
||||
for idx, entry in enumerate(entries):
|
||||
assert entry.version == 2
|
||||
assert not entry.options
|
||||
assert len(entry.subentries) == 1
|
||||
subentry = list(entry.subentries.values())[0]
|
||||
assert subentry.subentry_type == "conversation"
|
||||
assert subentry.data == TEST_OPTIONS
|
||||
assert subentry.title == f"Ollama {idx + 1}"
|
||||
|
||||
dev = device_registry.async_get_device(
|
||||
identifiers={(DOMAIN, list(entry.subentries.values())[0].subentry_id)}
|
||||
)
|
||||
assert dev is not None
|
||||
|
||||
|
||||
async def test_migration_from_v1_to_v2_with_same_urls(
|
||||
hass: HomeAssistant,
|
||||
device_registry: dr.DeviceRegistry,
|
||||
entity_registry: er.EntityRegistry,
|
||||
) -> None:
|
||||
"""Test migration from version 1 to version 2 with same URLs consolidates entries."""
|
||||
# Create two v1 config entries with the same URL
|
||||
mock_config_entry = MockConfigEntry(
|
||||
domain=DOMAIN,
|
||||
data={"url": "http://localhost:11434", "model": "llama3.2:latest"},
|
||||
options=TEST_OPTIONS,
|
||||
version=1,
|
||||
title="Ollama",
|
||||
)
|
||||
mock_config_entry.add_to_hass(hass)
|
||||
mock_config_entry_2 = MockConfigEntry(
|
||||
domain=DOMAIN,
|
||||
data={"url": "http://localhost:11434", "model": "llama3.2:latest"}, # Same URL
|
||||
options=TEST_OPTIONS,
|
||||
version=1,
|
||||
title="Ollama 2",
|
||||
)
|
||||
mock_config_entry_2.add_to_hass(hass)
|
||||
|
||||
device = device_registry.async_get_or_create(
|
||||
config_entry_id=mock_config_entry.entry_id,
|
||||
identifiers={(DOMAIN, mock_config_entry.entry_id)},
|
||||
name=mock_config_entry.title,
|
||||
manufacturer="Ollama",
|
||||
model="Ollama",
|
||||
entry_type=dr.DeviceEntryType.SERVICE,
|
||||
)
|
||||
entity_registry.async_get_or_create(
|
||||
"conversation",
|
||||
DOMAIN,
|
||||
mock_config_entry.entry_id,
|
||||
config_entry=mock_config_entry,
|
||||
device_id=device.id,
|
||||
suggested_object_id="ollama",
|
||||
)
|
||||
|
||||
device_2 = device_registry.async_get_or_create(
|
||||
config_entry_id=mock_config_entry_2.entry_id,
|
||||
identifiers={(DOMAIN, mock_config_entry_2.entry_id)},
|
||||
name=mock_config_entry_2.title,
|
||||
manufacturer="Ollama",
|
||||
model="Ollama",
|
||||
entry_type=dr.DeviceEntryType.SERVICE,
|
||||
)
|
||||
entity_registry.async_get_or_create(
|
||||
"conversation",
|
||||
DOMAIN,
|
||||
mock_config_entry_2.entry_id,
|
||||
config_entry=mock_config_entry_2,
|
||||
device_id=device_2.id,
|
||||
suggested_object_id="ollama_2",
|
||||
)
|
||||
|
||||
# Run migration
|
||||
with patch(
|
||||
"homeassistant.components.ollama.async_setup_entry",
|
||||
return_value=True,
|
||||
):
|
||||
await hass.config_entries.async_setup(mock_config_entry.entry_id)
|
||||
await hass.async_block_till_done()
|
||||
|
||||
# Should have only one entry left (consolidated)
|
||||
entries = hass.config_entries.async_entries(DOMAIN)
|
||||
assert len(entries) == 1
|
||||
|
||||
entry = entries[0]
|
||||
assert entry.version == 2
|
||||
assert not entry.options
|
||||
assert len(entry.subentries) == 2 # Two subentries from the two original entries
|
||||
|
||||
# Check both subentries exist with correct data
|
||||
subentries = list(entry.subentries.values())
|
||||
titles = [sub.title for sub in subentries]
|
||||
assert "Ollama" in titles
|
||||
assert "Ollama 2" in titles
|
||||
|
||||
for subentry in subentries:
|
||||
assert subentry.subentry_type == "conversation"
|
||||
assert subentry.data == TEST_OPTIONS
|
||||
|
||||
# Check devices were migrated correctly
|
||||
dev = device_registry.async_get_device(
|
||||
identifiers={(DOMAIN, subentry.subentry_id)}
|
||||
)
|
||||
assert dev is not None
|
||||
|
Reference in New Issue
Block a user