Correctly retrieve only loaded Google Generative AI config_entries (#139999)

* Correctly retrieve only loaded config_entries

* Ruff
This commit is contained in:
Ivan Lopez Hernandez 2025-03-06 13:17:33 -08:00 committed by GitHub
parent e78139edf1
commit 2aa584ce39
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 67 additions and 3 deletions

View File

@ -65,9 +65,9 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
prompt_parts = [call.data[CONF_PROMPT]]
config_entry: GoogleGenerativeAIConfigEntry = hass.config_entries.async_entries(
DOMAIN
)[0]
config_entry: GoogleGenerativeAIConfigEntry = (
hass.config_entries.async_loaded_entries(DOMAIN)[0]
)
client = config_entry.runtime_data

View File

@ -31,3 +31,18 @@
),
])
# ---
# name: test_load_entry_with_unloaded_entries
list([
tuple(
'',
tuple(
),
dict({
'contents': list([
'Write an opening speech for a Home Assistant release party',
]),
'model': 'models/gemini-2.0-flash',
}),
),
])
# ---

View File

@ -224,3 +224,52 @@ async def test_config_entry_error(
await hass.async_block_till_done()
assert mock_config_entry.state == state
assert any(mock_config_entry.async_get_active_flows(hass, {"reauth"})) == reauth
@pytest.mark.usefixtures("mock_init_component")
async def test_load_entry_with_unloaded_entries(
hass: HomeAssistant, snapshot: SnapshotAssertion
) -> None:
"""Test loading an entry with unloaded entries."""
config_entries = hass.config_entries.async_entries(
"google_generative_ai_conversation"
)
runtime_data = config_entries[0].runtime_data
await hass.config_entries.async_unload(config_entries[0].entry_id)
entry = MockConfigEntry(
domain="google_generative_ai_conversation",
title="Google Generative AI Conversation",
data={
"api_key": "bla",
},
state=ConfigEntryState.LOADED,
)
entry.runtime_data = runtime_data
entry.add_to_hass(hass)
stubbed_generated_content = (
"I'm thrilled to welcome you all to the release "
"party for the latest version of Home Assistant!"
)
with patch(
"google.genai.models.AsyncModels.generate_content",
return_value=Mock(
text=stubbed_generated_content,
prompt_feedback=None,
candidates=[Mock()],
),
) as mock_generate:
response = await hass.services.async_call(
"google_generative_ai_conversation",
"generate_content",
{"prompt": "Write an opening speech for a Home Assistant release party"},
blocking=True,
return_response=True,
)
assert response == {
"text": stubbed_generated_content,
}
assert [tuple(mock_call) for mock_call in mock_generate.mock_calls] == snapshot