jeans/core
jeans/core
1
0
Fork 0
mirror of https://github.com/home-assistant/core.git synced 2025-07-27 15:17:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Create Google Generative AI sub entries for an enabled entry (#148161)

Browse Source 
Co-authored-by: Erik Montnemery <erik@montnemery.com>
This commit is contained in:
Joost Lekkerkerker 2025-07-14 11:24:00 +02:00 committed by GitHub
parent 9f3d890e91
commit 334d5f09fb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 520 additions and 30 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

117
homeassistant/components/google_generative_ai_conversation/__init__.py
View File

@ -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,6 +351,52 @@ async def async_migrate_entry(
if entry.version == 2 and entry.minor_version == 2:
# Add AI Task subentry with default options
_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(
@ -324,10 +406,3 @@ async def async_migrate_entry(
unique_id=None,
),
)
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

2
homeassistant/components/google_generative_ai_conversation/config_flow.py
View File

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

427
tests/components/google_generative_ai_conversation/test_init.py
View File

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