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

* timeout

* fix

* tests
This commit is contained in:
tronikos 2024-05-24 18:31:02 -07:00 committed by GitHub
parent c9a79f6293
commit 5ca27f5d0c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 46 additions and 5 deletions

View File

@ -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

View File

@ -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",

View File

@ -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

View File

@ -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