diff --git a/homeassistant/components/google_generative_ai_conversation/__init__.py b/homeassistant/components/google_generative_ai_conversation/__init__.py index d1b8467955a..563d7d341f9 100644 --- a/homeassistant/components/google_generative_ai_conversation/__init__.py +++ b/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)) - except ClientError as err: - if err.reason == "API_KEY_INVALID": + async with timeout(5.0): + next(await hass.async_add_executor_job(partial(genai.list_models)), None) + 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 diff --git a/tests/components/google_generative_ai_conversation/conftest.py b/tests/components/google_generative_ai_conversation/conftest.py index 4dfa6379d73..8ab8020428e 100644 --- a/tests/components/google_generative_ai_conversation/conftest.py +++ b/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", diff --git a/tests/components/google_generative_ai_conversation/test_config_flow.py b/tests/components/google_generative_ai_conversation/test_config_flow.py index 460d74734ae..ba21407343e 100644 --- a/tests/components/google_generative_ai_conversation/test_config_flow.py +++ b/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 diff --git a/tests/components/google_generative_ai_conversation/test_init.py b/tests/components/google_generative_ai_conversation/test_init.py index 7dfa8bebfa5..a6a5fdf0b0e 100644 --- a/tests/components/google_generative_ai_conversation/test_init.py +++ b/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