From 2aa584ce39901e68339da802c631d1a89532a0c1 Mon Sep 17 00:00:00 2001
From: Ivan Lopez Hernandez <ivan.lh.94@outlook.com>
Date: Thu, 6 Mar 2025 13:17:33 -0800
Subject: [PATCH] Correctly retrieve only loaded Google Generative AI
config_entries (#139999)
* Correctly retrieve only loaded config_entries
* Ruff
---
.../__init__.py | 6 +--
.../snapshots/test_init.ambr | 15 ++++++
.../test_init.py | 49 +++++++++++++++++++
3 files changed, 67 insertions(+), 3 deletions(-)
diff --git a/homeassistant/components/google_generative_ai_conversation/__init__.py b/homeassistant/components/google_generative_ai_conversation/__init__.py
index 33e361d1433..6b10565e0b5 100644
--- a/homeassistant/components/google_generative_ai_conversation/__init__.py
+++ b/homeassistant/components/google_generative_ai_conversation/__init__.py
@@ -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
diff --git a/tests/components/google_generative_ai_conversation/snapshots/test_init.ambr b/tests/components/google_generative_ai_conversation/snapshots/test_init.ambr
index 8e6231cbffd..ce882adf6e6 100644
--- a/tests/components/google_generative_ai_conversation/snapshots/test_init.ambr
+++ b/tests/components/google_generative_ai_conversation/snapshots/test_init.ambr
@@ -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',
+ }),
+ ),
+ ])
+# ---
diff --git a/tests/components/google_generative_ai_conversation/test_init.py b/tests/components/google_generative_ai_conversation/test_init.py
index 0dad485812e..25533ffd46e 100644
--- a/tests/components/google_generative_ai_conversation/test_init.py
+++ b/tests/components/google_generative_ai_conversation/test_init.py
@@ -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
--
GitLab