From 334d5f09fb693343dff7094d018205ceb577dfdf Mon Sep 17 00:00:00 2001 From: Joost Lekkerkerker Date: Mon, 14 Jul 2025 11:24:00 +0200 Subject: [PATCH] Create Google Generative AI sub entries for an enabled entry (#148161) Co-authored-by: Erik Montnemery --- .../__init__.py | 121 ++++- .../config_flow.py | 2 +- .../test_init.py | 427 +++++++++++++++++- 3 files changed, 520 insertions(+), 30 deletions(-) diff --git a/homeassistant/components/google_generative_ai_conversation/__init__.py b/homeassistant/components/google_generative_ai_conversation/__init__.py index a3b87c05e5a..1ff9f355c06 100644 --- a/homeassistant/components/google_generative_ai_conversation/__init__.py +++ b/homeassistant/components/google_generative_ai_conversation/__init__.py @@ -195,11 +195,15 @@ async def async_update_options( 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] = {} + api_keys_entries: dict[str, tuple[ConfigEntry, bool]] = {} entity_registry = er.async_get(hass) device_registry = dr.async_get(hass) @@ -213,9 +217,14 @@ async def async_migrate_integration(hass: HomeAssistant) -> None: ) if entry.data[CONF_API_KEY] not in api_keys_entries: use_existing = True - api_keys_entries[entry.data[CONF_API_KEY]] = entry + all_disabled = all( + e.disabled_by is not None + for e in entries + if e.data[CONF_API_KEY] == entry.data[CONF_API_KEY] + ) + api_keys_entries[entry.data[CONF_API_KEY]] = (entry, all_disabled) - parent_entry = api_keys_entries[entry.data[CONF_API_KEY]] + parent_entry, all_disabled = api_keys_entries[entry.data[CONF_API_KEY]] hass.config_entries.async_add_subentry(parent_entry, subentry) if use_existing: @@ -228,25 +237,51 @@ async def async_migrate_integration(hass: HomeAssistant) -> None: unique_id=None, ), ) - 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, @@ -266,12 +301,13 @@ 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_TITLE, options={}, version=2, - minor_version=2, + minor_version=4, ) @@ -315,19 +351,58 @@ async def async_migrate_entry( if entry.version == 2 and entry.minor_version == 2: # Add AI Task subentry with default options - hass.config_entries.async_add_subentry( - entry, - ConfigSubentry( - data=MappingProxyType(RECOMMENDED_AI_TASK_OPTIONS), - 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=3) + if entry.version == 2 and entry.minor_version == 3: + # 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=4) + LOGGER.debug( "Migration to version %s:%s successful", entry.version, entry.minor_version ) return True + + +def _add_ai_task_subentry( + hass: HomeAssistant, entry: GoogleGenerativeAIConfigEntry +) -> None: + """Add AI Task subentry to the config entry.""" + hass.config_entries.async_add_subentry( + entry, + ConfigSubentry( + data=MappingProxyType(RECOMMENDED_AI_TASK_OPTIONS), + subentry_type="ai_task_data", + title=DEFAULT_AI_TASK_NAME, + unique_id=None, + ), + ) diff --git a/homeassistant/components/google_generative_ai_conversation/config_flow.py b/homeassistant/components/google_generative_ai_conversation/config_flow.py index a68ca09e76d..7d1429b110e 100644 --- a/homeassistant/components/google_generative_ai_conversation/config_flow.py +++ b/homeassistant/components/google_generative_ai_conversation/config_flow.py @@ -97,7 +97,7 @@ class GoogleGenerativeAIConfigFlow(ConfigFlow, domain=DOMAIN): """Handle a config flow for Google Generative AI Conversation.""" VERSION = 2 - MINOR_VERSION = 3 + MINOR_VERSION = 4 async def async_step_api( self, user_input: dict[str, Any] | None = None diff --git a/tests/components/google_generative_ai_conversation/test_init.py b/tests/components/google_generative_ai_conversation/test_init.py index 351293e7ac0..e154f9d33c9 100644 --- a/tests/components/google_generative_ai_conversation/test_init.py +++ b/tests/components/google_generative_ai_conversation/test_init.py @@ -1,5 +1,6 @@ """Tests for the Google Generative AI Conversation integration.""" +from typing import Any from unittest.mock import AsyncMock, Mock, mock_open, patch from google.genai.types import File, FileState @@ -17,11 +18,17 @@ from homeassistant.components.google_generative_ai_conversation.const import ( RECOMMENDED_CONVERSATION_OPTIONS, RECOMMENDED_TTS_OPTIONS, ) -from homeassistant.config_entries import ConfigEntryState, ConfigSubentryData +from homeassistant.config_entries import ( + ConfigEntryDisabler, + ConfigEntryState, + ConfigSubentryData, +) from homeassistant.const import CONF_API_KEY from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers import device_registry as dr, entity_registry as er +from homeassistant.helpers.device_registry import DeviceEntryDisabler +from homeassistant.helpers.entity_registry import RegistryEntryDisabler from . import API_ERROR_500, CLIENT_ERROR_API_KEY_INVALID @@ -479,7 +486,7 @@ async def test_migration_from_v1( assert len(entries) == 1 entry = entries[0] assert entry.version == 2 - assert entry.minor_version == 3 + assert entry.minor_version == 4 assert not entry.options assert entry.title == DEFAULT_TITLE assert len(entry.subentries) == 4 @@ -556,6 +563,223 @@ async def test_migration_from_v1( } +@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.google_generative_ai_conversation_2", + "device_disabled_by": None, + "entity_disabled_by": None, + "device": 1, + }, + { + "conversation_entity_id": "conversation.google_generative_ai_conversation", + "device_disabled_by": DeviceEntryDisabler.USER, + "entity_disabled_by": RegistryEntryDisabler.DEVICE, + "device": 0, + }, + ], + 1, + ), + ( + [None, ConfigEntryDisabler.USER], + None, + [ + { + "conversation_entity_id": "conversation.google_generative_ai_conversation", + "device_disabled_by": DeviceEntryDisabler.USER, + "entity_disabled_by": RegistryEntryDisabler.DEVICE, + "device": 0, + }, + { + "conversation_entity_id": "conversation.google_generative_ai_conversation_2", + "device_disabled_by": None, + "entity_disabled_by": None, + "device": 1, + }, + ], + 0, + ), + ( + [ConfigEntryDisabler.USER, ConfigEntryDisabler.USER], + ConfigEntryDisabler.USER, + [ + { + "conversation_entity_id": "conversation.google_generative_ai_conversation", + "device_disabled_by": DeviceEntryDisabler.CONFIG_ENTRY, + "entity_disabled_by": RegistryEntryDisabler.CONFIG_ENTRY, + "device": 0, + }, + { + "conversation_entity_id": "conversation.google_generative_ai_conversation_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 + options = { + "recommended": True, + "llm_hass_api": ["assist"], + "prompt": "You are a helpful assistant", + "chat_model": "models/gemini-2.0-flash", + } + mock_config_entry = MockConfigEntry( + domain=DOMAIN, + data={CONF_API_KEY: "1234"}, + options=options, + version=1, + title="Google Generative AI", + disabled_by=config_entry_disabled_by[0], + ) + mock_config_entry.add_to_hass(hass) + mock_config_entry_2 = MockConfigEntry( + domain=DOMAIN, + data={CONF_API_KEY: "1234"}, + options=options, + version=1, + title="Google Generative AI 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="Google", + model="Generative AI", + 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="google_generative_ai_conversation", + 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="Google", + model="Generative AI", + 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="google_generative_ai_conversation_2", + ) + + devices = [device_1, device_2] + + # Run migration + with patch( + "homeassistant.components.google_generative_ai_conversation.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 == 2 + assert entry.minor_version == 4 + assert not entry.options + assert entry.title == DEFAULT_TITLE + assert len(entry.subentries) == 4 + 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 == options + assert "Google Generative AI" in subentry.title + tts_subentries = [ + subentry + for subentry in entry.subentries.values() + if subentry.subentry_type == "tts" + ] + assert len(tts_subentries) == 1 + assert tts_subentries[0].data == RECOMMENDED_TTS_OPTIONS + assert tts_subentries[0].title == DEFAULT_TTS_NAME + 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 == RECOMMENDED_AI_TASK_OPTIONS + assert ai_task_subentries[0].title == DEFAULT_AI_TASK_NAME + + 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_v1_with_multiple_keys( hass: HomeAssistant, device_registry: dr.DeviceRegistry, @@ -633,7 +857,7 @@ async def test_migration_from_v1_with_multiple_keys( for entry in entries: assert entry.version == 2 - assert entry.minor_version == 3 + assert entry.minor_version == 4 assert not entry.options assert entry.title == DEFAULT_TITLE assert len(entry.subentries) == 3 @@ -736,7 +960,7 @@ async def test_migration_from_v1_with_same_keys( assert len(entries) == 1 entry = entries[0] assert entry.version == 2 - assert entry.minor_version == 3 + assert entry.minor_version == 4 assert not entry.options assert entry.title == DEFAULT_TITLE assert len(entry.subentries) == 4 @@ -957,7 +1181,7 @@ async def test_migration_from_v2_1( assert len(entries) == 1 entry = entries[0] assert entry.version == 2 - assert entry.minor_version == 3 + assert entry.minor_version == 4 assert not entry.options assert entry.title == DEFAULT_TITLE assert len(entry.subentries) == 4 @@ -1094,7 +1318,7 @@ async def test_migrate_entry_from_v2_2(hass: HomeAssistant) -> None: # Check version and subversion were updated assert entry.version == 2 - assert entry.minor_version == 3 + assert entry.minor_version == 4 # Check we now have conversation, tts and ai_task_data subentries assert len(entry.subentries) == 3 @@ -1123,3 +1347,194 @@ async def test_migrate_entry_from_v2_2(hass: HomeAssistant) -> None: assert tts_subentry is not None assert tts_subentry.title == DEFAULT_TTS_NAME assert tts_subentry.data == RECOMMENDED_TTS_OPTIONS + + +@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, + 4, + None, + DeviceEntryDisabler.USER, + RegistryEntryDisabler.DEVICE, + ), + ( + None, + DeviceEntryDisabler.USER, + RegistryEntryDisabler.DEVICE, + True, + 4, + None, + DeviceEntryDisabler.USER, + RegistryEntryDisabler.DEVICE, + ), + ( + None, + DeviceEntryDisabler.USER, + RegistryEntryDisabler.USER, + True, + 4, + None, + DeviceEntryDisabler.USER, + RegistryEntryDisabler.USER, + ), + ( + None, + None, + None, + True, + 4, + None, + None, + None, + ), + # Config entry disabled, migration does not run + ( + ConfigEntryDisabler.USER, + DeviceEntryDisabler.CONFIG_ENTRY, + RegistryEntryDisabler.CONFIG_ENTRY, + False, + 3, + ConfigEntryDisabler.USER, + DeviceEntryDisabler.CONFIG_ENTRY, + RegistryEntryDisabler.CONFIG_ENTRY, + ), + ( + ConfigEntryDisabler.USER, + DeviceEntryDisabler.USER, + RegistryEntryDisabler.DEVICE, + False, + 3, + ConfigEntryDisabler.USER, + DeviceEntryDisabler.USER, + RegistryEntryDisabler.DEVICE, + ), + ( + ConfigEntryDisabler.USER, + DeviceEntryDisabler.USER, + RegistryEntryDisabler.USER, + False, + 3, + ConfigEntryDisabler.USER, + DeviceEntryDisabler.USER, + RegistryEntryDisabler.USER, + ), + ( + ConfigEntryDisabler.USER, + None, + None, + False, + 3, + ConfigEntryDisabler.USER, + None, + None, + ), + ], +) +async def test_migrate_entry_from_v2_3( + 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 2.3.""" + # Create a v2.3 config entry with conversation and TTS subentries + conversation_subentry_id = "blabla" + mock_config_entry = MockConfigEntry( + domain=DOMAIN, + data={CONF_API_KEY: "test-api-key"}, + disabled_by=config_entry_disabled_by, + version=2, + minor_version=3, + subentries_data=[ + { + "data": RECOMMENDED_CONVERSATION_OPTIONS, + "subentry_id": conversation_subentry_id, + "subentry_type": "conversation", + "title": DEFAULT_CONVERSATION_NAME, + "unique_id": None, + }, + { + "data": RECOMMENDED_TTS_OPTIONS, + "subentry_type": "tts", + "title": DEFAULT_TTS_NAME, + "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="Google", + model="Generative AI", + 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="google_generative_ai_conversation", + ) + + # Verify initial state + assert mock_config_entry.version == 2 + assert mock_config_entry.minor_version == 3 + 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.google_generative_ai_conversation.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 == 2 + 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