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:
Paulus Schoutsen
2025-06-24 16:13:34 -04:00
committed by GitHub
parent 5a20ef3f3f
commit f735331699
8 changed files with 625 additions and 139 deletions

View File

@@ -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

View File

@@ -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"),
[

View File

@@ -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,
},
)

View File

@@ -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