Improve type hints in google_generative_ai tests (#121170)

This commit is contained in:
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.""" """Tests helpers."""
from collections.abc import Generator
from unittest.mock import patch from unittest.mock import patch
import pytest import pytest
@ -14,14 +15,14 @@ from tests.common import MockConfigEntry
@pytest.fixture @pytest.fixture
def mock_genai(): def mock_genai() -> Generator[None]:
"""Mock the genai call in async_setup_entry.""" """Mock the genai call in async_setup_entry."""
with patch("google.ai.generativelanguage_v1beta.ModelServiceAsyncClient.get_model"): with patch("google.ai.generativelanguage_v1beta.ModelServiceAsyncClient.get_model"):
yield yield
@pytest.fixture @pytest.fixture
def mock_config_entry(hass, mock_genai): def mock_config_entry(hass: HomeAssistant, mock_genai: None) -> MockConfigEntry:
"""Mock a config entry.""" """Mock a config entry."""
entry = MockConfigEntry( entry = MockConfigEntry(
domain="google_generative_ai_conversation", domain="google_generative_ai_conversation",
@ -35,7 +36,9 @@ def mock_config_entry(hass, mock_genai):
@pytest.fixture @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.""" """Mock a config entry with assist."""
hass.config_entries.async_update_entry( hass.config_entries.async_update_entry(
mock_config_entry, options={CONF_LLM_HASS_API: llm.LLM_API_ASSIST} 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 @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.""" """Initialize integration."""
assert await async_setup_component(hass, "google_generative_ai_conversation", {}) assert await async_setup_component(hass, "google_generative_ai_conversation", {})
await hass.async_block_till_done() 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( async def test_options_switching(
hass: HomeAssistant, hass: HomeAssistant,
mock_config_entry, mock_config_entry: MockConfigEntry,
mock_init_component,
mock_models, mock_models,
current_options, current_options,
new_options, new_options,

View File

@ -44,10 +44,10 @@ def freeze_the_time():
{CONF_LLM_HASS_API: llm.LLM_API_ASSIST}, {CONF_LLM_HASS_API: llm.LLM_API_ASSIST},
], ],
) )
@pytest.mark.usefixtures("mock_init_component")
async def test_default_prompt( async def test_default_prompt(
hass: HomeAssistant, hass: HomeAssistant,
mock_config_entry: MockConfigEntry, mock_config_entry: MockConfigEntry,
mock_init_component,
snapshot: SnapshotAssertion, snapshot: SnapshotAssertion,
agent_id: str | None, agent_id: str | None,
config_entry_options: {}, config_entry_options: {},
@ -102,10 +102,10 @@ async def test_default_prompt(
("model_name", "supports_system_instruction"), ("model_name", "supports_system_instruction"),
[("models/gemini-1.5-pro", True), ("models/gemini-1.0-pro", False)], [("models/gemini-1.5-pro", True), ("models/gemini-1.0-pro", False)],
) )
@pytest.mark.usefixtures("mock_init_component")
async def test_chat_history( async def test_chat_history(
hass: HomeAssistant, hass: HomeAssistant,
mock_config_entry: MockConfigEntry, mock_config_entry: MockConfigEntry,
mock_init_component,
model_name: str, model_name: str,
supports_system_instruction: bool, supports_system_instruction: bool,
snapshot: SnapshotAssertion, snapshot: SnapshotAssertion,
@ -167,11 +167,11 @@ async def test_chat_history(
@patch( @patch(
"homeassistant.components.google_generative_ai_conversation.conversation.llm.AssistAPI._async_get_tools" "homeassistant.components.google_generative_ai_conversation.conversation.llm.AssistAPI._async_get_tools"
) )
@pytest.mark.usefixtures("mock_init_component")
async def test_function_call( async def test_function_call(
mock_get_tools, mock_get_tools,
hass: HomeAssistant, hass: HomeAssistant,
mock_config_entry_with_assist: MockConfigEntry, mock_config_entry_with_assist: MockConfigEntry,
mock_init_component,
snapshot: SnapshotAssertion, snapshot: SnapshotAssertion,
) -> None: ) -> None:
"""Test function calling.""" """Test function calling."""
@ -277,11 +277,11 @@ async def test_function_call(
@patch( @patch(
"homeassistant.components.google_generative_ai_conversation.conversation.llm.AssistAPI._async_get_tools" "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( async def test_function_call_without_parameters(
mock_get_tools, mock_get_tools,
hass: HomeAssistant, hass: HomeAssistant,
mock_config_entry_with_assist: MockConfigEntry, mock_config_entry_with_assist: MockConfigEntry,
mock_init_component,
snapshot: SnapshotAssertion, snapshot: SnapshotAssertion,
) -> None: ) -> None:
"""Test function calling without parameters.""" """Test function calling without parameters."""
@ -358,11 +358,11 @@ async def test_function_call_without_parameters(
@patch( @patch(
"homeassistant.components.google_generative_ai_conversation.conversation.llm.AssistAPI._async_get_tools" "homeassistant.components.google_generative_ai_conversation.conversation.llm.AssistAPI._async_get_tools"
) )
@pytest.mark.usefixtures("mock_init_component")
async def test_function_exception( async def test_function_exception(
mock_get_tools, mock_get_tools,
hass: HomeAssistant, hass: HomeAssistant,
mock_config_entry_with_assist: MockConfigEntry, mock_config_entry_with_assist: MockConfigEntry,
mock_init_component,
) -> None: ) -> None:
"""Test exception in function calling.""" """Test exception in function calling."""
agent_id = mock_config_entry_with_assist.entry_id 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( async def test_error_handling(
hass: HomeAssistant, mock_config_entry: MockConfigEntry, mock_init_component hass: HomeAssistant, mock_config_entry: MockConfigEntry
) -> None: ) -> None:
"""Test that client errors are caught.""" """Test that client errors are caught."""
with patch("google.generativeai.GenerativeModel") as mock_model: 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( async def test_blocked_response(
hass: HomeAssistant, mock_config_entry: MockConfigEntry, mock_init_component hass: HomeAssistant, mock_config_entry: MockConfigEntry
) -> None: ) -> None:
"""Test blocked response.""" """Test blocked response."""
with patch("google.generativeai.GenerativeModel") as mock_model: 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( async def test_empty_response(
hass: HomeAssistant, mock_config_entry: MockConfigEntry, mock_init_component hass: HomeAssistant, mock_config_entry: MockConfigEntry
) -> None: ) -> None:
"""Test empty response.""" """Test empty response."""
with patch("google.generativeai.GenerativeModel") as mock_model: 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( async def test_invalid_llm_api(
hass: HomeAssistant, hass: HomeAssistant, mock_config_entry: MockConfigEntry
mock_config_entry: MockConfigEntry,
mock_init_component,
) -> None: ) -> None:
"""Test handling of invalid llm api.""" """Test handling of invalid llm api."""
hass.config_entries.async_update_entry( 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"] 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( async def test_conversation_agent(
hass: HomeAssistant, hass: HomeAssistant, mock_config_entry: MockConfigEntry
mock_config_entry: MockConfigEntry,
mock_init_component,
) -> None: ) -> None:
"""Test GoogleGenerativeAIAgent.""" """Test GoogleGenerativeAIAgent."""
agent = conversation.get_agent_manager(hass).async_get_agent( 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 from tests.common import MockConfigEntry
@pytest.mark.usefixtures("mock_init_component")
async def test_generate_content_service_without_images( async def test_generate_content_service_without_images(
hass: HomeAssistant, hass: HomeAssistant, snapshot: SnapshotAssertion
mock_config_entry: MockConfigEntry,
mock_init_component,
snapshot: SnapshotAssertion,
) -> None: ) -> None:
"""Test generate content service.""" """Test generate content service."""
stubbed_generated_content = ( 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 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( async def test_generate_content_service_with_image(
hass: HomeAssistant, hass: HomeAssistant, snapshot: SnapshotAssertion
mock_config_entry: MockConfigEntry,
mock_init_component,
snapshot: SnapshotAssertion,
) -> None: ) -> None:
"""Test generate content service.""" """Test generate content service."""
stubbed_generated_content = ( 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( async def test_generate_content_service_with_image_not_allowed_path(
hass: HomeAssistant, hass: HomeAssistant,
mock_config_entry: MockConfigEntry,
mock_init_component,
snapshot: SnapshotAssertion,
) -> None: ) -> None:
"""Test generate content service with an image in a not allowed path.""" """Test generate content service with an image in a not allowed path."""
with ( 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( async def test_generate_content_service_with_image_not_exists(
hass: HomeAssistant, hass: HomeAssistant,
mock_config_entry: MockConfigEntry,
mock_init_component,
snapshot: SnapshotAssertion,
) -> None: ) -> None:
"""Test generate content service with an image that does not exist.""" """Test generate content service with an image that does not exist."""
with ( with (
@ -192,12 +184,8 @@ async def test_generate_content_service_with_image_not_exists(
) )
async def test_generate_content_service_with_non_image( @pytest.mark.usefixtures("mock_init_component")
hass: HomeAssistant, async def test_generate_content_service_with_non_image(hass: HomeAssistant) -> None:
mock_config_entry: MockConfigEntry,
mock_init_component,
snapshot: SnapshotAssertion,
) -> None:
"""Test generate content service with a non image.""" """Test generate content service with a non image."""
with ( with (
patch("pathlib.Path.exists", return_value=True), patch("pathlib.Path.exists", return_value=True),