mirror of
https://github.com/home-assistant/core.git
synced 2025-07-24 05:37:44 +00:00
Google Generative AI: add timeout to ensure we don't block HA startup (#118066)
* timeout * fix * tests
This commit is contained in:
parent
c9a79f6293
commit
5ca27f5d0c
@ -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
|
||||
|
@ -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",
|
||||
|
@ -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
|
||||
|
||||
|
@ -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
|
||||
|
Loading…
x
Reference in New Issue
Block a user