From 43e4223a8efb9a1739b7d8a16733381735bae401 Mon Sep 17 00:00:00 2001 From: epenet <6771947+epenet@users.noreply.github.com> Date: Thu, 4 Jul 2024 11:09:34 +0200 Subject: [PATCH] Improve type hints in google_generative_ai tests (#121170) --- .../conftest.py | 13 ++++++--- .../test_config_flow.py | 4 +-- .../test_conversation.py | 29 ++++++++++--------- .../test_init.py | 28 +++++------------- 4 files changed, 34 insertions(+), 40 deletions(-) diff --git a/tests/components/google_generative_ai_conversation/conftest.py b/tests/components/google_generative_ai_conversation/conftest.py index 1761516e4f5..28c21a9b791 100644 --- a/tests/components/google_generative_ai_conversation/conftest.py +++ b/tests/components/google_generative_ai_conversation/conftest.py @@ -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() 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 c835a4d3b13..d4992c732e1 100644 --- a/tests/components/google_generative_ai_conversation/test_config_flow.py +++ b/tests/components/google_generative_ai_conversation/test_config_flow.py @@ -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, diff --git a/tests/components/google_generative_ai_conversation/test_conversation.py b/tests/components/google_generative_ai_conversation/test_conversation.py index 1e45c79a3b6..dc713f09454 100644 --- a/tests/components/google_generative_ai_conversation/test_conversation.py +++ b/tests/components/google_generative_ai_conversation/test_conversation.py @@ -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( diff --git a/tests/components/google_generative_ai_conversation/test_init.py b/tests/components/google_generative_ai_conversation/test_init.py index 3f860c78da1..4875323d094 100644 --- a/tests/components/google_generative_ai_conversation/test_init.py +++ b/tests/components/google_generative_ai_conversation/test_init.py @@ -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),