Migrate Google Gen AI to use subentries (#147281)

* Migrate Google Gen AI to use subentries

* Add reconfig successful msg

* Address comments

* Do not allow addin subentry when not loaded

* Let HA do the migration

* Use config_entries.async_setup

* Remove fallback name on base entity

* Fix

* Fix

* Fix device name assignment in entity and tts modules

* Fix tests

---------

Co-authored-by: Joostlek <joostlek@outlook.com>
This commit is contained in:
Paulus Schoutsen
2025-06-23 20:59:32 -04:00
committed by GitHub
parent 6641cb3799
commit 56f4039ac2
12 changed files with 599 additions and 122 deletions

View File

@@ -5,8 +5,9 @@ from unittest.mock import AsyncMock, Mock, patch
import pytest
from homeassistant.components.google_generative_ai_conversation.entity import (
from homeassistant.components.google_generative_ai_conversation.const import (
CONF_USE_GOOGLE_SEARCH_TOOL,
DEFAULT_CONVERSATION_NAME,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_LLM_HASS_API
@@ -26,6 +27,15 @@ def mock_config_entry(hass: HomeAssistant) -> MockConfigEntry:
data={
"api_key": "bla",
},
version=2,
subentries_data=[
{
"data": {},
"subentry_type": "conversation",
"title": DEFAULT_CONVERSATION_NAME,
"unique_id": None,
}
],
)
entry.runtime_data = Mock()
entry.add_to_hass(hass)
@@ -38,8 +48,10 @@ async def mock_config_entry_with_assist(
) -> MockConfigEntry:
"""Mock a config entry with assist."""
with patch("google.genai.models.AsyncModels.get"):
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},
)
await hass.async_block_till_done()
return mock_config_entry
@@ -51,9 +63,10 @@ async def mock_config_entry_with_google_search(
) -> MockConfigEntry:
"""Mock a config entry with assist."""
with patch("google.genai.models.AsyncModels.get"):
hass.config_entries.async_update_entry(
hass.config_entries.async_update_subentry(
mock_config_entry,
options={
next(iter(mock_config_entry.subentries.values())),
data={
CONF_LLM_HASS_API: llm.LLM_API_ASSIST,
CONF_USE_GOOGLE_SEARCH_TOOL: True,
},

View File

@@ -22,6 +22,7 @@ from homeassistant.components.google_generative_ai_conversation.const import (
CONF_TOP_K,
CONF_TOP_P,
CONF_USE_GOOGLE_SEARCH_TOOL,
DEFAULT_CONVERSATION_NAME,
DOMAIN,
RECOMMENDED_CHAT_MODEL,
RECOMMENDED_HARM_BLOCK_THRESHOLD,
@@ -30,7 +31,7 @@ from homeassistant.components.google_generative_ai_conversation.const import (
RECOMMENDED_TOP_P,
RECOMMENDED_USE_GOOGLE_SEARCH_TOOL,
)
from homeassistant.const import CONF_LLM_HASS_API
from homeassistant.const import CONF_API_KEY, CONF_LLM_HASS_API, CONF_NAME
from homeassistant.core import HomeAssistant
from homeassistant.data_entry_flow import FlowResultType
@@ -110,10 +111,100 @@ async def test_form(hass: HomeAssistant) -> None:
assert result2["data"] == {
"api_key": "bla",
}
assert result2["options"] == RECOMMENDED_OPTIONS
assert result2["options"] == {}
assert result2["subentries"] == [
{
"subentry_type": "conversation",
"data": RECOMMENDED_OPTIONS,
"title": DEFAULT_CONVERSATION_NAME,
"unique_id": None,
}
]
assert len(mock_setup_entry.mock_calls) == 1
async def test_duplicate_entry(hass: HomeAssistant) -> None:
"""Test we get the form."""
MockConfigEntry(
domain=DOMAIN,
data={CONF_API_KEY: "bla"},
).add_to_hass(hass)
result = await hass.config_entries.flow.async_init(
DOMAIN, context={"source": config_entries.SOURCE_USER}
)
assert result["type"] is FlowResultType.FORM
assert not result["errors"]
result = await hass.config_entries.flow.async_configure(
result["flow_id"],
{
CONF_API_KEY: "bla",
},
)
assert result["type"] is FlowResultType.ABORT
assert result["reason"] == "already_configured"
async def test_creating_conversation_subentry(
hass: HomeAssistant,
mock_init_component: None,
mock_config_entry: MockConfigEntry,
) -> None:
"""Test creating a conversation subentry."""
with patch(
"google.genai.models.AsyncModels.list",
return_value=get_models_pager(),
):
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.FORM
assert result["step_id"] == "set_options"
assert not result["errors"]
with patch(
"google.genai.models.AsyncModels.list",
return_value=get_models_pager(),
):
result2 = await hass.config_entries.subentries.async_configure(
result["flow_id"],
{CONF_NAME: "Mock name", **RECOMMENDED_OPTIONS},
)
await hass.async_block_till_done()
assert result2["type"] is FlowResultType.CREATE_ENTRY
assert result2["title"] == "Mock name"
processed_options = RECOMMENDED_OPTIONS.copy()
processed_options[CONF_PROMPT] = processed_options[CONF_PROMPT].strip()
assert result2["data"] == processed_options
async def test_creating_conversation_subentry_not_loaded(
hass: HomeAssistant,
mock_init_component: None,
mock_config_entry: MockConfigEntry,
) -> None:
"""Test creating a conversation subentry."""
await hass.config_entries.async_unload(mock_config_entry.entry_id)
with patch(
"google.genai.models.AsyncModels.list",
return_value=get_models_pager(),
):
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"
def will_options_be_rendered_again(current_options, new_options) -> bool:
"""Determine if options will be rendered again."""
return current_options.get(CONF_RECOMMENDED) != new_options.get(CONF_RECOMMENDED)
@@ -283,7 +374,7 @@ def will_options_be_rendered_again(current_options, new_options) -> bool:
],
)
@pytest.mark.usefixtures("mock_init_component")
async def test_options_switching(
async def test_subentry_options_switching(
hass: HomeAssistant,
mock_config_entry: MockConfigEntry,
current_options,
@@ -292,17 +383,18 @@ async def test_options_switching(
errors,
) -> None:
"""Test the options form."""
subentry = next(iter(mock_config_entry.subentries.values()))
with patch("google.genai.models.AsyncModels.get"):
hass.config_entries.async_update_entry(
mock_config_entry, options=current_options
hass.config_entries.async_update_subentry(
mock_config_entry, subentry, data=current_options
)
await hass.async_block_till_done()
with patch(
"google.genai.models.AsyncModels.list",
return_value=get_models_pager(),
):
options_flow = await hass.config_entries.options.async_init(
mock_config_entry.entry_id
options_flow = await mock_config_entry.start_subentry_reconfigure_flow(
hass, subentry.subentry_id
)
if will_options_be_rendered_again(current_options, new_options):
retry_options = {
@@ -313,7 +405,7 @@ async def test_options_switching(
"google.genai.models.AsyncModels.list",
return_value=get_models_pager(),
):
options_flow = await hass.config_entries.options.async_configure(
options_flow = await hass.config_entries.subentries.async_configure(
options_flow["flow_id"],
retry_options,
)
@@ -321,14 +413,15 @@ async def test_options_switching(
"google.genai.models.AsyncModels.list",
return_value=get_models_pager(),
):
options = await hass.config_entries.options.async_configure(
options = await hass.config_entries.subentries.async_configure(
options_flow["flow_id"],
new_options,
)
await hass.async_block_till_done()
await hass.async_block_till_done()
if errors is None:
assert options["type"] is FlowResultType.CREATE_ENTRY
assert options["data"] == expected_options
assert options["type"] is FlowResultType.ABORT
assert options["reason"] == "reconfigure_successful"
assert subentry.data == expected_options
else:
assert options["type"] is FlowResultType.FORM
@@ -375,7 +468,10 @@ async def test_reauth_flow(hass: HomeAssistant) -> None:
"""Test the reauth flow."""
hass.config.components.add("google_generative_ai_conversation")
mock_config_entry = MockConfigEntry(
domain=DOMAIN, state=config_entries.ConfigEntryState.LOADED, title="Gemini"
domain=DOMAIN,
state=config_entries.ConfigEntryState.LOADED,
title="Gemini",
version=2,
)
mock_config_entry.add_to_hass(hass)
mock_config_entry.async_start_reauth(hass)

View File

@@ -64,7 +64,7 @@ async def test_error_handling(
"hello",
None,
Context(),
agent_id="conversation.google_generative_ai_conversation",
agent_id="conversation.google_ai_conversation",
)
assert result.response.response_type == intent.IntentResponseType.ERROR, result
assert result.response.error_code == "unknown", result
@@ -82,7 +82,7 @@ async def test_function_call(
mock_send_message_stream: AsyncMock,
) -> None:
"""Test function calling."""
agent_id = "conversation.google_generative_ai_conversation"
agent_id = "conversation.google_ai_conversation"
context = Context()
messages = [
@@ -212,7 +212,7 @@ async def test_google_search_tool_is_sent(
mock_send_message_stream: AsyncMock,
) -> None:
"""Test if the Google Search tool is sent to the model."""
agent_id = "conversation.google_generative_ai_conversation"
agent_id = "conversation.google_ai_conversation"
context = Context()
messages = [
@@ -278,7 +278,7 @@ async def test_blocked_response(
mock_send_message_stream: AsyncMock,
) -> None:
"""Test blocked response."""
agent_id = "conversation.google_generative_ai_conversation"
agent_id = "conversation.google_ai_conversation"
context = Context()
messages = [
@@ -328,7 +328,7 @@ async def test_empty_response(
) -> None:
"""Test empty response."""
agent_id = "conversation.google_generative_ai_conversation"
agent_id = "conversation.google_ai_conversation"
context = Context()
messages = [
@@ -371,7 +371,7 @@ async def test_none_response(
mock_send_message_stream: AsyncMock,
) -> None:
"""Test None response."""
agent_id = "conversation.google_generative_ai_conversation"
agent_id = "conversation.google_ai_conversation"
context = Context()
messages = [
@@ -403,10 +403,12 @@ async def test_converse_error(
hass: HomeAssistant, mock_config_entry: MockConfigEntry
) -> None:
"""Test handling ChatLog raising ConverseError."""
subentry = next(iter(mock_config_entry.subentries.values()))
with patch("google.genai.models.AsyncModels.get"):
hass.config_entries.async_update_entry(
hass.config_entries.async_update_subentry(
mock_config_entry,
options={**mock_config_entry.options, CONF_LLM_HASS_API: "invalid_llm_api"},
next(iter(mock_config_entry.subentries.values())),
data={**subentry.data, CONF_LLM_HASS_API: "invalid_llm_api"},
)
await hass.async_block_till_done()
@@ -415,7 +417,7 @@ async def test_converse_error(
"hello",
None,
Context(),
agent_id="conversation.google_generative_ai_conversation",
agent_id="conversation.google_ai_conversation",
)
assert result.response.response_type == intent.IntentResponseType.ERROR, result
@@ -593,7 +595,7 @@ async def test_empty_content_in_chat_history(
mock_send_message_stream: AsyncMock,
) -> None:
"""Tests that in case of an empty entry in the chat history the google API will receive an injected space sign instead."""
agent_id = "conversation.google_generative_ai_conversation"
agent_id = "conversation.google_ai_conversation"
context = Context()
messages = [
@@ -648,7 +650,7 @@ async def test_history_always_user_first_turn(
) -> None:
"""Test that the user is always first in the chat history."""
agent_id = "conversation.google_generative_ai_conversation"
agent_id = "conversation.google_ai_conversation"
context = Context()
messages = [
@@ -674,7 +676,7 @@ async def test_history_always_user_first_turn(
mock_chat_log.async_add_assistant_content_without_tools(
conversation.AssistantContent(
agent_id="conversation.google_generative_ai_conversation",
agent_id="conversation.google_ai_conversation",
content="Garage door left open, do you want to close it?",
)
)

View File

@@ -7,9 +7,12 @@ import pytest
from requests.exceptions import Timeout
from syrupy.assertion import SnapshotAssertion
from homeassistant.components.google_generative_ai_conversation.const import DOMAIN
from homeassistant.config_entries import ConfigEntryState
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 . import API_ERROR_500, CLIENT_ERROR_API_KEY_INVALID
@@ -387,3 +390,214 @@ async def test_load_entry_with_unloaded_entries(
"text": stubbed_generated_content,
}
assert [tuple(mock_call) for mock_call in mock_generate.mock_calls] == snapshot
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
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",
)
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",
)
mock_config_entry_2.add_to_hass(hass)
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,
)
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",
)
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",
)
# 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.version == 2
assert not entry.options
assert len(entry.subentries) == 2
for subentry in entry.subentries.values():
assert subentry.subentry_type == "conversation"
assert subentry.data == options
assert "Google Generative AI" in subentry.title
subentry = list(entry.subentries.values())[0]
entity = entity_registry.async_get("conversation.google_generative_ai_conversation")
assert entity.unique_id == subentry.subentry_id
assert entity.config_subentry_id == subentry.subentry_id
assert entity.config_entry_id == entry.entry_id
assert not device_registry.async_get_device(
identifiers={(DOMAIN, mock_config_entry.entry_id)}
)
assert (
device := device_registry.async_get_device(
identifiers={(DOMAIN, subentry.subentry_id)}
)
)
assert device.identifiers == {(DOMAIN, subentry.subentry_id)}
assert device.id == device_1.id
subentry = list(entry.subentries.values())[1]
entity = entity_registry.async_get(
"conversation.google_generative_ai_conversation_2"
)
assert entity.unique_id == subentry.subentry_id
assert entity.config_subentry_id == subentry.subentry_id
assert entity.config_entry_id == entry.entry_id
assert not device_registry.async_get_device(
identifiers={(DOMAIN, mock_config_entry_2.entry_id)}
)
assert (
device := device_registry.async_get_device(
identifiers={(DOMAIN, subentry.subentry_id)}
)
)
assert device.identifiers == {(DOMAIN, subentry.subentry_id)}
assert device.id == device_2.id
async def test_migration_from_v1_to_v2_with_multiple_keys(
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
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",
)
mock_config_entry.add_to_hass(hass)
mock_config_entry_2 = MockConfigEntry(
domain=DOMAIN,
data={CONF_API_KEY: "12345"},
options=options,
version=1,
title="Google Generative AI 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="Google",
model="Generative AI",
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="google_generative_ai_conversation",
)
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",
)
# 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) == 2
for entry in 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 == options
assert "Google Generative AI" in subentry.title
dev = device_registry.async_get_device(
identifiers={(DOMAIN, list(entry.subentries.values())[0].subentry_id)}
)
assert dev is not None

View File

@@ -122,7 +122,9 @@ async def mock_setup(hass: HomeAssistant, config: dict[str, Any]) -> None:
async def mock_config_entry_setup(hass: HomeAssistant, config: dict[str, Any]) -> None:
"""Mock config entry setup."""
default_config = {tts.CONF_LANG: "en-US"}
config_entry = MockConfigEntry(domain=DOMAIN, data=default_config | config)
config_entry = MockConfigEntry(
domain=DOMAIN, data=default_config | config, version=2
)
client_mock = Mock()
client_mock.models.get = None