Improve type hints in google_generative_ai tests (#121170)

pull/121037/head^2
epenet 2024-07-04 11:09:34 +02:00 committed by GitHub
parent 869f24df49
commit 43e4223a8e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 34 additions and 40 deletions

View File

@ -1,5 +1,6 @@
"""Tests helpers."""
from collections.abc import Generator
from unittest.mock import patch
import pytest
@ -14,14 +15,14 @@ from tests.common import MockConfigEntry
@pytest.fixture
def mock_genai():
def mock_genai() -> Generator[None]:
"""Mock the genai call in async_setup_entry."""
with patch("google.ai.generativelanguage_v1beta.ModelServiceAsyncClient.get_model"):
yield
@pytest.fixture
def mock_config_entry(hass, mock_genai):
def mock_config_entry(hass: HomeAssistant, mock_genai: None) -> MockConfigEntry:
"""Mock a config entry."""
entry = MockConfigEntry(
domain="google_generative_ai_conversation",
@ -35,7 +36,9 @@ def mock_config_entry(hass, mock_genai):
@pytest.fixture
def mock_config_entry_with_assist(hass, mock_config_entry):
def mock_config_entry_with_assist(
hass: HomeAssistant, mock_config_entry: MockConfigEntry
) -> MockConfigEntry:
"""Mock a config entry with assist."""
hass.config_entries.async_update_entry(
mock_config_entry, options={CONF_LLM_HASS_API: llm.LLM_API_ASSIST}
@ -44,7 +47,9 @@ def mock_config_entry_with_assist(hass, mock_config_entry):
@pytest.fixture
async def mock_init_component(hass: HomeAssistant, mock_config_entry: ConfigEntry):
async def mock_init_component(
hass: HomeAssistant, mock_config_entry: ConfigEntry
) -> None:
"""Initialize integration."""
assert await async_setup_component(hass, "google_generative_ai_conversation", {})
await hass.async_block_till_done()

View File

@ -154,10 +154,10 @@ async def test_form(hass: HomeAssistant) -> None:
),
],
)
@pytest.mark.usefixtures("mock_init_component")
async def test_options_switching(
hass: HomeAssistant,
mock_config_entry,
mock_init_component,
mock_config_entry: MockConfigEntry,
mock_models,
current_options,
new_options,

View File

@ -44,10 +44,10 @@ def freeze_the_time():
{CONF_LLM_HASS_API: llm.LLM_API_ASSIST},
],
)
@pytest.mark.usefixtures("mock_init_component")
async def test_default_prompt(
hass: HomeAssistant,
mock_config_entry: MockConfigEntry,
mock_init_component,
snapshot: SnapshotAssertion,
agent_id: str | None,
config_entry_options: {},
@ -102,10 +102,10 @@ async def test_default_prompt(
("model_name", "supports_system_instruction"),
[("models/gemini-1.5-pro", True), ("models/gemini-1.0-pro", False)],
)
@pytest.mark.usefixtures("mock_init_component")
async def test_chat_history(
hass: HomeAssistant,
mock_config_entry: MockConfigEntry,
mock_init_component,
model_name: str,
supports_system_instruction: bool,
snapshot: SnapshotAssertion,
@ -167,11 +167,11 @@ async def test_chat_history(
@patch(
"homeassistant.components.google_generative_ai_conversation.conversation.llm.AssistAPI._async_get_tools"
)
@pytest.mark.usefixtures("mock_init_component")
async def test_function_call(
mock_get_tools,
hass: HomeAssistant,
mock_config_entry_with_assist: MockConfigEntry,
mock_init_component,
snapshot: SnapshotAssertion,
) -> None:
"""Test function calling."""
@ -277,11 +277,11 @@ async def test_function_call(
@patch(
"homeassistant.components.google_generative_ai_conversation.conversation.llm.AssistAPI._async_get_tools"
)
@pytest.mark.usefixtures("mock_init_component")
async def test_function_call_without_parameters(
mock_get_tools,
hass: HomeAssistant,
mock_config_entry_with_assist: MockConfigEntry,
mock_init_component,
snapshot: SnapshotAssertion,
) -> None:
"""Test function calling without parameters."""
@ -358,11 +358,11 @@ async def test_function_call_without_parameters(
@patch(
"homeassistant.components.google_generative_ai_conversation.conversation.llm.AssistAPI._async_get_tools"
)
@pytest.mark.usefixtures("mock_init_component")
async def test_function_exception(
mock_get_tools,
hass: HomeAssistant,
mock_config_entry_with_assist: MockConfigEntry,
mock_init_component,
) -> None:
"""Test exception in function calling."""
agent_id = mock_config_entry_with_assist.entry_id
@ -440,8 +440,9 @@ async def test_function_exception(
)
@pytest.mark.usefixtures("mock_init_component")
async def test_error_handling(
hass: HomeAssistant, mock_config_entry: MockConfigEntry, mock_init_component
hass: HomeAssistant, mock_config_entry: MockConfigEntry
) -> None:
"""Test that client errors are caught."""
with patch("google.generativeai.GenerativeModel") as mock_model:
@ -459,8 +460,9 @@ async def test_error_handling(
)
@pytest.mark.usefixtures("mock_init_component")
async def test_blocked_response(
hass: HomeAssistant, mock_config_entry: MockConfigEntry, mock_init_component
hass: HomeAssistant, mock_config_entry: MockConfigEntry
) -> None:
"""Test blocked response."""
with patch("google.generativeai.GenerativeModel") as mock_model:
@ -480,8 +482,9 @@ async def test_blocked_response(
)
@pytest.mark.usefixtures("mock_init_component")
async def test_empty_response(
hass: HomeAssistant, mock_config_entry: MockConfigEntry, mock_init_component
hass: HomeAssistant, mock_config_entry: MockConfigEntry
) -> None:
"""Test empty response."""
with patch("google.generativeai.GenerativeModel") as mock_model:
@ -501,10 +504,9 @@ async def test_empty_response(
)
@pytest.mark.usefixtures("mock_init_component")
async def test_invalid_llm_api(
hass: HomeAssistant,
mock_config_entry: MockConfigEntry,
mock_init_component,
hass: HomeAssistant, mock_config_entry: MockConfigEntry
) -> None:
"""Test handling of invalid llm api."""
hass.config_entries.async_update_entry(
@ -593,10 +595,9 @@ async def test_template_variables(
assert "The user id is 12345." in mock_model.mock_calls[0][2]["system_instruction"]
@pytest.mark.usefixtures("mock_init_component")
async def test_conversation_agent(
hass: HomeAssistant,
mock_config_entry: MockConfigEntry,
mock_init_component,
hass: HomeAssistant, mock_config_entry: MockConfigEntry
) -> None:
"""Test GoogleGenerativeAIAgent."""
agent = conversation.get_agent_manager(hass).async_get_agent(

View File

@ -14,11 +14,9 @@ from homeassistant.exceptions import HomeAssistantError
from tests.common import MockConfigEntry
@pytest.mark.usefixtures("mock_init_component")
async def test_generate_content_service_without_images(
hass: HomeAssistant,
mock_config_entry: MockConfigEntry,
mock_init_component,
snapshot: SnapshotAssertion,
hass: HomeAssistant, snapshot: SnapshotAssertion
) -> None:
"""Test generate content service."""
stubbed_generated_content = (
@ -46,11 +44,9 @@ async def test_generate_content_service_without_images(
assert [tuple(mock_call) for mock_call in mock_model.mock_calls] == snapshot
@pytest.mark.usefixtures("mock_init_component")
async def test_generate_content_service_with_image(
hass: HomeAssistant,
mock_config_entry: MockConfigEntry,
mock_init_component,
snapshot: SnapshotAssertion,
hass: HomeAssistant, snapshot: SnapshotAssertion
) -> None:
"""Test generate content service."""
stubbed_generated_content = (
@ -134,11 +130,9 @@ async def test_generate_content_response_has_empty_parts(
)
@pytest.mark.usefixtures("mock_init_component")
async def test_generate_content_service_with_image_not_allowed_path(
hass: HomeAssistant,
mock_config_entry: MockConfigEntry,
mock_init_component,
snapshot: SnapshotAssertion,
) -> None:
"""Test generate content service with an image in a not allowed path."""
with (
@ -165,11 +159,9 @@ async def test_generate_content_service_with_image_not_allowed_path(
)
@pytest.mark.usefixtures("mock_init_component")
async def test_generate_content_service_with_image_not_exists(
hass: HomeAssistant,
mock_config_entry: MockConfigEntry,
mock_init_component,
snapshot: SnapshotAssertion,
) -> None:
"""Test generate content service with an image that does not exist."""
with (
@ -192,12 +184,8 @@ async def test_generate_content_service_with_image_not_exists(
)
async def test_generate_content_service_with_non_image(
hass: HomeAssistant,
mock_config_entry: MockConfigEntry,
mock_init_component,
snapshot: SnapshotAssertion,
) -> None:
@pytest.mark.usefixtures("mock_init_component")
async def test_generate_content_service_with_non_image(hass: HomeAssistant) -> None:
"""Test generate content service with a non image."""
with (
patch("pathlib.Path.exists", return_value=True),