Google Generative AI: add timeout to ensure we don't block HA startup (#118066)

* timeout

* fix

* tests

homeassistant/components/google_generative_ai_conversation/__init__.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from asyncio import timeout
 from functools import partial
 import mimetypes
 from pathlib import Path
@@ -100,9 +101,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     genai.configure(api_key=entry.data[CONF_API_KEY])
 
     try:
-        await hass.async_add_executor_job(partial(genai.list_models))
+        async with timeout(5.0):
-    except ClientError as err:
+            next(await hass.async_add_executor_job(partial(genai.list_models)), None)
-        if err.reason == "API_KEY_INVALID":
+    except (ClientError, TimeoutError) as err:
+        if isinstance(err, ClientError) and err.reason == "API_KEY_INVALID":
             LOGGER.error("Invalid API key: %s", err)
             return False
         raise ConfigEntryNotReady(err) from err

tests/components/google_generative_ai_conversation/conftest.py

@@ -14,7 +14,17 @@ from tests.common import MockConfigEntry
 
 
 @pytest.fixture
-def mock_config_entry(hass):
+def mock_genai():
+    """Mock the genai call in async_setup_entry."""
+    with patch(
+        "homeassistant.components.google_generative_ai_conversation.genai.list_models",
+        return_value=iter([]),
+    ):
+        yield
+
+
+@pytest.fixture
+def mock_config_entry(hass, mock_genai):
     """Mock a config entry."""
     entry = MockConfigEntry(
         domain="google_generative_ai_conversation",

tests/components/google_generative_ai_conversation/test_config_flow.py

@@ -45,7 +45,7 @@ def mock_models():
     model_10_pro.name = "models/gemini-pro"
     with patch(
         "homeassistant.components.google_generative_ai_conversation.config_flow.genai.list_models",
-        return_value=[model_15_flash, model_10_pro],
+        return_value=iter([model_15_flash, model_10_pro]),
     ):
         yield
 

tests/components/google_generative_ai_conversation/test_init.py

@@ -6,6 +6,7 @@ from google.api_core.exceptions import ClientError
 import pytest
 from syrupy.assertion import SnapshotAssertion
 
+from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 
@@ -217,3 +218,31 @@ async def test_generate_content_service_with_non_image(
             blocking=True,
             return_response=True,
         )
+
+
+async def test_config_entry_not_ready(
+    hass: HomeAssistant, mock_config_entry: MockConfigEntry
+) -> None:
+    """Test configuration entry not ready."""
+    with patch(
+        "homeassistant.components.google_generative_ai_conversation.genai.list_models",
+        side_effect=ClientError("error"),
+    ):
+        mock_config_entry.add_to_hass(hass)
+        await hass.config_entries.async_setup(mock_config_entry.entry_id)
+        await hass.async_block_till_done()
+        assert mock_config_entry.state is ConfigEntryState.SETUP_RETRY
+
+
+async def test_config_entry_setup_error(
+    hass: HomeAssistant, mock_config_entry: MockConfigEntry
+) -> None:
+    """Test configuration entry setup error."""
+    with patch(
+        "homeassistant.components.google_generative_ai_conversation.genai.list_models",
+        side_effect=ClientError("error", error_info="API_KEY_INVALID"),
+    ):
+        mock_config_entry.add_to_hass(hass)
+        await hass.config_entries.async_setup(mock_config_entry.entry_id)
+        await hass.async_block_till_done()
+        assert mock_config_entry.state is ConfigEntryState.SETUP_ERROR