From a54f0adf74dd5db836ef420ac08c0e221f360f07 Mon Sep 17 00:00:00 2001 From: Joost Lekkerkerker Date: Wed, 6 Aug 2025 14:27:36 +0200 Subject: [PATCH] Enable disabled Ollama config entries after entry migration (#150105) --- homeassistant/components/ollama/__init__.py | 147 +++++-- .../components/ollama/config_flow.py | 2 +- tests/components/ollama/test_init.py | 412 +++++++++++++++++- 3 files changed, 516 insertions(+), 45 deletions(-) diff --git a/homeassistant/components/ollama/__init__.py b/homeassistant/components/ollama/__init__.py index e16550c1e94..091e58dbe7f 100644 --- a/homeassistant/components/ollama/__init__.py +++ b/homeassistant/components/ollama/__init__.py @@ -92,11 +92,15 @@ async def async_update_options(hass: HomeAssistant, entry: OllamaConfigEntry) -> async def async_migrate_integration(hass: HomeAssistant) -> None: """Migrate integration entry structure.""" - entries = hass.config_entries.async_entries(DOMAIN) + # Make sure we get enabled config entries first + entries = sorted( + hass.config_entries.async_entries(DOMAIN), + key=lambda e: e.disabled_by is not None, + ) if not any(entry.version == 1 for entry in entries): return - api_keys_entries: dict[str, ConfigEntry] = {} + url_entries: dict[str, tuple[ConfigEntry, bool]] = {} entity_registry = er.async_get(hass) device_registry = dr.async_get(hass) @@ -112,33 +116,64 @@ async def async_migrate_integration(hass: HomeAssistant) -> None: title=entry.title, unique_id=None, ) - if entry.data[CONF_URL] not in api_keys_entries: + if entry.data[CONF_URL] not in url_entries: use_existing = True - api_keys_entries[entry.data[CONF_URL]] = entry + all_disabled = all( + e.disabled_by is not None + for e in entries + if e.data[CONF_URL] == entry.data[CONF_URL] + ) + url_entries[entry.data[CONF_URL]] = (entry, all_disabled) - parent_entry = api_keys_entries[entry.data[CONF_URL]] + parent_entry, all_disabled = url_entries[entry.data[CONF_URL]] hass.config_entries.async_add_subentry(parent_entry, subentry) - conversation_entity = entity_registry.async_get_entity_id( + conversation_entity_id = entity_registry.async_get_entity_id( "conversation", DOMAIN, entry.entry_id, ) - if conversation_entity is not None: - entity_registry.async_update_entity( - conversation_entity, - config_entry_id=parent_entry.entry_id, - config_subentry_id=subentry.subentry_id, - new_unique_id=subentry.subentry_id, - ) - device = device_registry.async_get_device( identifiers={(DOMAIN, entry.entry_id)} ) + + if conversation_entity_id is not None: + conversation_entity_entry = entity_registry.entities[conversation_entity_id] + entity_disabled_by = conversation_entity_entry.disabled_by + if ( + entity_disabled_by is er.RegistryEntryDisabler.CONFIG_ENTRY + and not all_disabled + ): + # Device and entity registries don't update the disabled_by flag + # when moving a device or entity from one config entry to another, + # so we need to do it manually. + entity_disabled_by = ( + er.RegistryEntryDisabler.DEVICE + if device + else er.RegistryEntryDisabler.USER + ) + entity_registry.async_update_entity( + conversation_entity_id, + config_entry_id=parent_entry.entry_id, + config_subentry_id=subentry.subentry_id, + disabled_by=entity_disabled_by, + new_unique_id=subentry.subentry_id, + ) + if device is not None: + # Device and entity registries don't update the disabled_by flag when + # moving a device or entity from one config entry to another, so we + # need to do it manually. + device_disabled_by = device.disabled_by + if ( + device.disabled_by is dr.DeviceEntryDisabler.CONFIG_ENTRY + and not all_disabled + ): + device_disabled_by = dr.DeviceEntryDisabler.USER device_registry.async_update_device( device.id, + disabled_by=device_disabled_by, new_identifiers={(DOMAIN, subentry.subentry_id)}, add_config_subentry_id=subentry.subentry_id, add_config_entry_id=parent_entry.entry_id, @@ -158,6 +193,7 @@ async def async_migrate_integration(hass: HomeAssistant) -> None: if not use_existing: await hass.config_entries.async_remove(entry.entry_id) else: + _add_ai_task_subentry(hass, entry) hass.config_entries.async_update_entry( entry, title=DEFAULT_NAME, @@ -165,7 +201,7 @@ async def async_migrate_integration(hass: HomeAssistant) -> None: data={CONF_URL: entry.data[CONF_URL]}, options={}, version=3, - minor_version=1, + minor_version=3, ) @@ -211,32 +247,69 @@ async def async_migrate_entry(hass: HomeAssistant, entry: OllamaConfigEntry) -> ) if entry.version == 3 and entry.minor_version == 1: - # Add AI Task subentry with default options. We can only create a new - # subentry if we can find an existing model in the entry. The model - # was removed in the previous migration step, so we need to - # check the subentries for an existing model. - existing_model = next( - iter( - model - for subentry in entry.subentries.values() - if (model := subentry.data.get(CONF_MODEL)) is not None - ), - None, - ) - if existing_model: - hass.config_entries.async_add_subentry( - entry, - ConfigSubentry( - data=MappingProxyType({CONF_MODEL: existing_model}), - subentry_type="ai_task_data", - title=DEFAULT_AI_TASK_NAME, - unique_id=None, - ), - ) + _add_ai_task_subentry(hass, entry) hass.config_entries.async_update_entry(entry, minor_version=2) + if entry.version == 3 and entry.minor_version == 2: + # Fix migration where the disabled_by flag was not set correctly. + # We can currently only correct this for enabled config entries, + # because migration does not run for disabled config entries. This + # is asserted in tests, and if that behavior is changed, we should + # correct also disabled config entries. + device_registry = dr.async_get(hass) + entity_registry = er.async_get(hass) + devices = dr.async_entries_for_config_entry(device_registry, entry.entry_id) + entity_entries = er.async_entries_for_config_entry( + entity_registry, entry.entry_id + ) + if entry.disabled_by is None: + # If the config entry is not disabled, we need to set the disabled_by + # flag on devices to USER, and on entities to DEVICE, if they are set + # to CONFIG_ENTRY. + for device in devices: + if device.disabled_by is not dr.DeviceEntryDisabler.CONFIG_ENTRY: + continue + device_registry.async_update_device( + device.id, + disabled_by=dr.DeviceEntryDisabler.USER, + ) + for entity in entity_entries: + if entity.disabled_by is not er.RegistryEntryDisabler.CONFIG_ENTRY: + continue + entity_registry.async_update_entity( + entity.entity_id, + disabled_by=er.RegistryEntryDisabler.DEVICE, + ) + hass.config_entries.async_update_entry(entry, minor_version=3) + _LOGGER.debug( "Migration to version %s:%s successful", entry.version, entry.minor_version ) return True + + +def _add_ai_task_subentry(hass: HomeAssistant, entry: OllamaConfigEntry) -> None: + """Add AI Task subentry to the config entry.""" + # Add AI Task subentry with default options. We can only create a new + # subentry if we can find an existing model in the entry. The model + # was removed in the previous migration step, so we need to + # check the subentries for an existing model. + existing_model = next( + iter( + model + for subentry in entry.subentries.values() + if (model := subentry.data.get(CONF_MODEL)) is not None + ), + None, + ) + if existing_model: + hass.config_entries.async_add_subentry( + entry, + ConfigSubentry( + data=MappingProxyType({CONF_MODEL: existing_model}), + subentry_type="ai_task_data", + title=DEFAULT_AI_TASK_NAME, + unique_id=None, + ), + ) diff --git a/homeassistant/components/ollama/config_flow.py b/homeassistant/components/ollama/config_flow.py index cca917f6c29..68deb00d205 100644 --- a/homeassistant/components/ollama/config_flow.py +++ b/homeassistant/components/ollama/config_flow.py @@ -76,7 +76,7 @@ class OllamaConfigFlow(ConfigFlow, domain=DOMAIN): """Handle a config flow for Ollama.""" VERSION = 3 - MINOR_VERSION = 2 + MINOR_VERSION = 3 def __init__(self) -> None: """Initialize config flow.""" diff --git a/tests/components/ollama/test_init.py b/tests/components/ollama/test_init.py index 1db57302704..766de8a7d6d 100644 --- a/tests/components/ollama/test_init.py +++ b/tests/components/ollama/test_init.py @@ -1,5 +1,6 @@ """Tests for the Ollama integration.""" +from typing import Any from unittest.mock import patch from httpx import ConnectError @@ -7,9 +8,12 @@ import pytest from homeassistant.components import ollama from homeassistant.components.ollama.const import DOMAIN -from homeassistant.config_entries import ConfigSubentryData +from homeassistant.config_entries import ConfigEntryDisabler, ConfigSubentryData +from homeassistant.const import CONF_URL from homeassistant.core import HomeAssistant from homeassistant.helpers import device_registry as dr, entity_registry as er, llm +from homeassistant.helpers.device_registry import DeviceEntryDisabler +from homeassistant.helpers.entity_registry import RegistryEntryDisabler from homeassistant.setup import async_setup_component from . import TEST_OPTIONS @@ -96,7 +100,7 @@ async def test_migration_from_v1( await hass.async_block_till_done() assert mock_config_entry.version == 3 - assert mock_config_entry.minor_version == 2 + assert mock_config_entry.minor_version == 3 # After migration, parent entry should only have URL assert mock_config_entry.data == {ollama.CONF_URL: "http://localhost:11434"} assert mock_config_entry.options == {} @@ -223,7 +227,7 @@ async def test_migration_from_v1_with_multiple_urls( for idx, entry in enumerate(entries): assert entry.version == 3 - assert entry.minor_version == 2 + assert entry.minor_version == 3 assert not entry.options assert len(entry.subentries) == 2 @@ -332,7 +336,7 @@ async def test_migration_from_v1_with_same_urls( entry = entries[0] assert entry.version == 3 - assert entry.minor_version == 2 + assert entry.minor_version == 3 assert not entry.options # Two conversation subentries from the two original entries and 1 aitask subentry assert len(entry.subentries) == 3 @@ -365,6 +369,209 @@ async def test_migration_from_v1_with_same_urls( } +@pytest.mark.parametrize( + ( + "config_entry_disabled_by", + "merged_config_entry_disabled_by", + "conversation_subentry_data", + "main_config_entry", + ), + [ + ( + [ConfigEntryDisabler.USER, None], + None, + [ + { + "conversation_entity_id": "conversation.ollama_2", + "device_disabled_by": None, + "entity_disabled_by": None, + "device": 1, + }, + { + "conversation_entity_id": "conversation.ollama", + "device_disabled_by": DeviceEntryDisabler.USER, + "entity_disabled_by": RegistryEntryDisabler.DEVICE, + "device": 0, + }, + ], + 1, + ), + ( + [None, ConfigEntryDisabler.USER], + None, + [ + { + "conversation_entity_id": "conversation.ollama", + "device_disabled_by": DeviceEntryDisabler.USER, + "entity_disabled_by": RegistryEntryDisabler.DEVICE, + "device": 0, + }, + { + "conversation_entity_id": "conversation.ollama_2", + "device_disabled_by": None, + "entity_disabled_by": None, + "device": 1, + }, + ], + 0, + ), + ( + [ConfigEntryDisabler.USER, ConfigEntryDisabler.USER], + ConfigEntryDisabler.USER, + [ + { + "conversation_entity_id": "conversation.ollama", + "device_disabled_by": DeviceEntryDisabler.CONFIG_ENTRY, + "entity_disabled_by": RegistryEntryDisabler.CONFIG_ENTRY, + "device": 0, + }, + { + "conversation_entity_id": "conversation.ollama_2", + "device_disabled_by": None, + "entity_disabled_by": None, + "device": 1, + }, + ], + 0, + ), + ], +) +async def test_migration_from_v1_disabled( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, + config_entry_disabled_by: list[ConfigEntryDisabler | None], + merged_config_entry_disabled_by: ConfigEntryDisabler | None, + conversation_subentry_data: list[dict[str, Any]], + main_config_entry: int, +) -> None: + """Test migration where the config entries are disabled.""" + # Create a v1 config entry with conversation options and an entity + mock_config_entry = MockConfigEntry( + domain=DOMAIN, + data={"url": "http://localhost:11434", "model": "llama3.2:latest"}, + options=V1_TEST_OPTIONS, + version=1, + title="Ollama", + disabled_by=config_entry_disabled_by[0], + ) + mock_config_entry.add_to_hass(hass) + mock_config_entry_2 = MockConfigEntry( + domain=DOMAIN, + data={"url": "http://localhost:11434", "model": "llama3.2:latest"}, + options=V1_TEST_OPTIONS, + version=1, + title="Ollama 2", + disabled_by=config_entry_disabled_by[1], + ) + mock_config_entry_2.add_to_hass(hass) + mock_config_entries = [mock_config_entry, mock_config_entry_2] + + device_1 = 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, + disabled_by=DeviceEntryDisabler.CONFIG_ENTRY, + ) + entity_registry.async_get_or_create( + "conversation", + DOMAIN, + mock_config_entry.entry_id, + config_entry=mock_config_entry, + device_id=device_1.id, + suggested_object_id="ollama", + disabled_by=RegistryEntryDisabler.CONFIG_ENTRY, + ) + + 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", + ) + + devices = [device_1, device_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) == 1 + entry = entries[0] + assert entry.disabled_by is merged_config_entry_disabled_by + assert entry.version == 3 + assert entry.minor_version == 3 + assert not entry.options + assert entry.title == "Ollama" + assert len(entry.subentries) == 3 + conversation_subentries = [ + subentry + for subentry in entry.subentries.values() + if subentry.subentry_type == "conversation" + ] + assert len(conversation_subentries) == 2 + for subentry in conversation_subentries: + assert subentry.subentry_type == "conversation" + assert subentry.data == {"model": "llama3.2:latest", **V1_TEST_OPTIONS} + assert "Ollama" in subentry.title + ai_task_subentries = [ + subentry + for subentry in entry.subentries.values() + if subentry.subentry_type == "ai_task_data" + ] + assert len(ai_task_subentries) == 1 + assert ai_task_subentries[0].data == {"model": "llama3.2:latest"} + assert ai_task_subentries[0].title == "Ollama AI Task" + + assert not device_registry.async_get_device( + identifiers={(DOMAIN, mock_config_entry.entry_id)} + ) + assert not device_registry.async_get_device( + identifiers={(DOMAIN, mock_config_entry_2.entry_id)} + ) + + for idx, subentry in enumerate(conversation_subentries): + subentry_data = conversation_subentry_data[idx] + entity = entity_registry.async_get(subentry_data["conversation_entity_id"]) + assert entity.unique_id == subentry.subentry_id + assert entity.config_subentry_id == subentry.subentry_id + assert entity.config_entry_id == entry.entry_id + assert entity.disabled_by is subentry_data["entity_disabled_by"] + + assert ( + device := device_registry.async_get_device( + identifiers={(DOMAIN, subentry.subentry_id)} + ) + ) + assert device.identifiers == {(DOMAIN, subentry.subentry_id)} + assert device.id == devices[subentry_data["device"]].id + assert device.config_entries == { + mock_config_entries[main_config_entry].entry_id + } + assert device.config_entries_subentries == { + mock_config_entries[main_config_entry].entry_id: {subentry.subentry_id} + } + assert device.disabled_by is subentry_data["device_disabled_by"] + + async def test_migration_from_v2_1( hass: HomeAssistant, device_registry: dr.DeviceRegistry, @@ -457,7 +664,7 @@ async def test_migration_from_v2_1( assert len(entries) == 1 entry = entries[0] assert entry.version == 3 - assert entry.minor_version == 2 + assert entry.minor_version == 3 assert not entry.options assert entry.title == "Ollama" assert len(entry.subentries) == 3 @@ -546,7 +753,7 @@ async def test_migration_from_v2_2(hass: HomeAssistant) -> None: # Check migration to v3.1 assert mock_config_entry.version == 3 - assert mock_config_entry.minor_version == 2 + assert mock_config_entry.minor_version == 3 # Check that model was moved from main data to subentry assert mock_config_entry.data == {ollama.CONF_URL: "http://localhost:11434"} @@ -584,6 +791,197 @@ async def test_migration_from_v3_1_without_subentry(hass: HomeAssistant) -> None await hass.config_entries.async_setup(mock_config_entry.entry_id) assert mock_config_entry.version == 3 - assert mock_config_entry.minor_version == 2 + assert mock_config_entry.minor_version == 3 assert next(iter(mock_config_entry.subentries.values()), None) is None + + +@pytest.mark.parametrize( + ( + "config_entry_disabled_by", + "device_disabled_by", + "entity_disabled_by", + "setup_result", + "minor_version_after_migration", + "config_entry_disabled_by_after_migration", + "device_disabled_by_after_migration", + "entity_disabled_by_after_migration", + ), + [ + # Config entry not disabled, update device and entity disabled by config entry + ( + None, + DeviceEntryDisabler.CONFIG_ENTRY, + RegistryEntryDisabler.CONFIG_ENTRY, + True, + 3, + None, + DeviceEntryDisabler.USER, + RegistryEntryDisabler.DEVICE, + ), + ( + None, + DeviceEntryDisabler.USER, + RegistryEntryDisabler.DEVICE, + True, + 3, + None, + DeviceEntryDisabler.USER, + RegistryEntryDisabler.DEVICE, + ), + ( + None, + DeviceEntryDisabler.USER, + RegistryEntryDisabler.USER, + True, + 3, + None, + DeviceEntryDisabler.USER, + RegistryEntryDisabler.USER, + ), + ( + None, + None, + None, + True, + 3, + None, + None, + None, + ), + # Config entry disabled, migration does not run + ( + ConfigEntryDisabler.USER, + DeviceEntryDisabler.CONFIG_ENTRY, + RegistryEntryDisabler.CONFIG_ENTRY, + False, + 2, + ConfigEntryDisabler.USER, + DeviceEntryDisabler.CONFIG_ENTRY, + RegistryEntryDisabler.CONFIG_ENTRY, + ), + ( + ConfigEntryDisabler.USER, + DeviceEntryDisabler.USER, + RegistryEntryDisabler.DEVICE, + False, + 2, + ConfigEntryDisabler.USER, + DeviceEntryDisabler.USER, + RegistryEntryDisabler.DEVICE, + ), + ( + ConfigEntryDisabler.USER, + DeviceEntryDisabler.USER, + RegistryEntryDisabler.USER, + False, + 2, + ConfigEntryDisabler.USER, + DeviceEntryDisabler.USER, + RegistryEntryDisabler.USER, + ), + ( + ConfigEntryDisabler.USER, + None, + None, + False, + 2, + ConfigEntryDisabler.USER, + None, + None, + ), + ], +) +async def test_migrate_entry_from_v3_2( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, + config_entry_disabled_by: ConfigEntryDisabler | None, + device_disabled_by: DeviceEntryDisabler | None, + entity_disabled_by: RegistryEntryDisabler | None, + setup_result: bool, + minor_version_after_migration: int, + config_entry_disabled_by_after_migration: ConfigEntryDisabler | None, + device_disabled_by_after_migration: ConfigEntryDisabler | None, + entity_disabled_by_after_migration: RegistryEntryDisabler | None, +) -> None: + """Test migration from version 3.2.""" + # Create a v3.2 config entry with conversation subentries + conversation_subentry_id = "blabla" + mock_config_entry = MockConfigEntry( + domain=DOMAIN, + data={CONF_URL: "http://localhost:11434"}, + disabled_by=config_entry_disabled_by, + version=3, + minor_version=2, + subentries_data=[ + { + "data": V1_TEST_OPTIONS, + "subentry_id": conversation_subentry_id, + "subentry_type": "conversation", + "title": "Ollama", + "unique_id": None, + }, + { + "data": {"model": "llama3.2:latest"}, + "subentry_type": "ai_task_data", + "title": "Ollama AI Task", + "unique_id": None, + }, + ], + ) + mock_config_entry.add_to_hass(hass) + + conversation_device = device_registry.async_get_or_create( + config_entry_id=mock_config_entry.entry_id, + config_subentry_id=conversation_subentry_id, + disabled_by=device_disabled_by, + identifiers={(DOMAIN, mock_config_entry.entry_id)}, + name=mock_config_entry.title, + manufacturer="Ollama", + model="Ollama", + entry_type=dr.DeviceEntryType.SERVICE, + ) + conversation_entity = entity_registry.async_get_or_create( + "conversation", + DOMAIN, + mock_config_entry.entry_id, + config_entry=mock_config_entry, + config_subentry_id=conversation_subentry_id, + disabled_by=entity_disabled_by, + device_id=conversation_device.id, + suggested_object_id="ollama", + ) + + # Verify initial state + assert mock_config_entry.version == 3 + assert mock_config_entry.minor_version == 2 + assert len(mock_config_entry.subentries) == 2 + assert mock_config_entry.disabled_by == config_entry_disabled_by + assert conversation_device.disabled_by == device_disabled_by + assert conversation_entity.disabled_by == entity_disabled_by + + # Run setup to trigger migration + with patch( + "homeassistant.components.ollama.async_setup_entry", + return_value=True, + ): + result = await hass.config_entries.async_setup(mock_config_entry.entry_id) + assert result is setup_result + await hass.async_block_till_done() + + # Verify migration completed + entries = hass.config_entries.async_entries(DOMAIN) + assert len(entries) == 1 + entry = entries[0] + + # Check version and subversion were updated + assert entry.version == 3 + assert entry.minor_version == minor_version_after_migration + + # Check the disabled_by flag on config entry, device and entity are as expected + conversation_device = device_registry.async_get(conversation_device.id) + conversation_entity = entity_registry.async_get(conversation_entity.entity_id) + assert mock_config_entry.disabled_by == config_entry_disabled_by_after_migration + assert conversation_device.disabled_by == device_disabled_by_after_migration + assert conversation_entity.disabled_by == entity_disabled_by_after_migration