Improve type hints in azure_event_hub tests (#119047)

Co-authored-by: Marc Mueller <30130371+cdce8p@users.noreply.github.com>

tests/components/azure_event_hub/conftest.py

@@ -3,10 +3,12 @@
 from dataclasses import dataclass
 from datetime import timedelta
 import logging
-from unittest.mock import MagicMock, patch
+from typing import Any
+from unittest.mock import AsyncMock, MagicMock, patch
 
 from azure.eventhub.aio import EventHubProducerClient
 import pytest
+from typing_extensions import AsyncGenerator, Generator
 
 from homeassistant.components.azure_event_hub.const import (
     CONF_FILTER,
@@ -15,6 +17,7 @@ from homeassistant.components.azure_event_hub.const import (
 )
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.const import STATE_ON
+from homeassistant.core import HomeAssistant
 from homeassistant.setup import async_setup_component
 from homeassistant.util.dt import utcnow
 
@@ -27,20 +30,25 @@ _LOGGER = logging.getLogger(__name__)
 
 # fixtures for both init and config flow tests
 @pytest.fixture(autouse=True, name="mock_get_eventhub_properties")
-def mock_get_eventhub_properties_fixture():
+def mock_get_eventhub_properties_fixture() -> Generator[AsyncMock]:
     """Mock azure event hub properties, used to test the connection."""
     with patch(f"{PRODUCER_PATH}.get_eventhub_properties") as get_eventhub_properties:
         yield get_eventhub_properties
 
 
 @pytest.fixture(name="filter_schema")
-def mock_filter_schema():
+def mock_filter_schema() -> dict[str, Any]:
     """Return an empty filter."""
     return {}
 
 
 @pytest.fixture(name="entry")
-async def mock_entry_fixture(hass, filter_schema, mock_create_batch, mock_send_batch):
+async def mock_entry_fixture(
+    hass: HomeAssistant,
+    filter_schema: dict[str, Any],
+    mock_create_batch: MagicMock,
+    mock_send_batch: AsyncMock,
+) -> AsyncGenerator[MockConfigEntry]:
     """Create the setup in HA."""
     entry = MockConfigEntry(
         domain=DOMAIN,
@@ -68,7 +76,9 @@ async def mock_entry_fixture(hass, filter_schema, mock_create_batch, mock_send_b
 
 # fixtures for init tests
 @pytest.fixture(name="entry_with_one_event")
-async def mock_entry_with_one_event(hass, entry):
+def mock_entry_with_one_event(
+    hass: HomeAssistant, entry: MockConfigEntry
+) -> MockConfigEntry:
     """Use the entry and add a single test event to the queue."""
     assert entry.state is ConfigEntryState.LOADED
     hass.states.async_set("sensor.test", STATE_ON)
@@ -84,14 +94,16 @@ class FilterTest:
 
 
 @pytest.fixture(name="mock_send_batch")
-def mock_send_batch_fixture():
+def mock_send_batch_fixture() -> Generator[AsyncMock]:
     """Mock send_batch."""
     with patch(f"{PRODUCER_PATH}.send_batch") as mock_send_batch:
         yield mock_send_batch
 
 
 @pytest.fixture(autouse=True, name="mock_client")
-def mock_client_fixture(mock_send_batch):
+def mock_client_fixture(
+    mock_send_batch: AsyncMock,
+) -> Generator[tuple[AsyncMock, AsyncMock]]:
     """Mock the azure event hub producer client."""
     with patch(f"{PRODUCER_PATH}.close") as mock_close:
         yield (
@@ -101,7 +113,7 @@ def mock_client_fixture(mock_send_batch):
 
 
 @pytest.fixture(name="mock_create_batch")
-def mock_create_batch_fixture():
+def mock_create_batch_fixture() -> Generator[MagicMock]:
     """Mock batch creator and return mocked batch object."""
     mock_batch = MagicMock()
     with patch(f"{PRODUCER_PATH}.create_batch", return_value=mock_batch):
@@ -110,7 +122,7 @@ def mock_create_batch_fixture():
 
 # fixtures for config flow tests
 @pytest.fixture(name="mock_from_connection_string")
-def mock_from_connection_string_fixture():
+def mock_from_connection_string_fixture() -> Generator[MagicMock]:
     """Mock AEH from connection string creation."""
     mock_aeh = MagicMock(spec=EventHubProducerClient)
     mock_aeh.__aenter__.return_value = mock_aeh
@@ -122,7 +134,7 @@ def mock_from_connection_string_fixture():
 
 
 @pytest.fixture
-def mock_setup_entry():
+def mock_setup_entry() -> Generator[AsyncMock]:
     """Mock the setup entry call, used for config flow tests."""
     with patch(
         f"{AZURE_EVENT_HUB_PATH}.async_setup_entry", return_value=True

tests/components/azure_event_hub/test_config_flow.py

@@ -1,7 +1,8 @@
 """Test the AEH config flow."""
 
 import logging
-from unittest.mock import AsyncMock
+from typing import Any
+from unittest.mock import AsyncMock, MagicMock
 
 from azure.eventhub.exceptions import EventHubError
 import pytest
@@ -43,14 +44,14 @@ pytestmark = pytest.mark.usefixtures("mock_setup_entry")
     ],
     ids=["connection_string", "sas"],
 )
+@pytest.mark.usefixtures("mock_from_connection_string")
 async def test_form(
     hass: HomeAssistant,
     mock_setup_entry: AsyncMock,
-    mock_from_connection_string,
-    step1_config,
-    step_id,
-    step2_config,
-    data_config,
+    step1_config: dict[str, Any],
+    step_id: str,
+    step2_config: dict[str, str],
+    data_config: dict[str, str],
 ) -> None:
     """Test we get the form."""
     result = await hass.config_entries.flow.async_init(
@@ -101,7 +102,7 @@ async def test_import(hass: HomeAssistant, mock_setup_entry: AsyncMock) -> None:
     [config_entries.SOURCE_USER, config_entries.SOURCE_IMPORT],
     ids=["user", "import"],
 )
-async def test_single_instance(hass: HomeAssistant, source) -> None:
+async def test_single_instance(hass: HomeAssistant, source: str) -> None:
     """Test uniqueness of username."""
     entry = MockConfigEntry(
         domain=DOMAIN,
@@ -126,9 +127,9 @@ async def test_single_instance(hass: HomeAssistant, source) -> None:
 )
 async def test_connection_error_sas(
     hass: HomeAssistant,
-    mock_get_eventhub_properties,
-    side_effect,
-    error_message,
+    mock_get_eventhub_properties: AsyncMock,
+    side_effect: Exception,
+    error_message: str,
 ) -> None:
     """Test we handle connection errors."""
     result = await hass.config_entries.flow.async_init(
@@ -155,9 +156,9 @@ async def test_connection_error_sas(
 )
 async def test_connection_error_cs(
     hass: HomeAssistant,
-    mock_from_connection_string,
-    side_effect,
-    error_message,
+    mock_from_connection_string: MagicMock,
+    side_effect: Exception,
+    error_message: str,
 ) -> None:
     """Test we handle connection errors."""
     result = await hass.config_entries.flow.async_init(
@@ -178,7 +179,7 @@ async def test_connection_error_cs(
     assert result2["errors"] == {"base": error_message}
 
 
-async def test_options_flow(hass: HomeAssistant, entry) -> None:
+async def test_options_flow(hass: HomeAssistant, entry: MockConfigEntry) -> None:
     """Test options flow."""
     result = await hass.config_entries.options.async_init(entry.entry_id)
 

tests/components/azure_event_hub/test_init.py

@@ -2,7 +2,7 @@
 
 from datetime import timedelta
 import logging
-from unittest.mock import patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 from azure.eventhub.exceptions import EventHubError
 import pytest
@@ -60,7 +60,9 @@ async def test_filter_only_config(hass: HomeAssistant) -> None:
     assert await async_setup_component(hass, DOMAIN, config)
 
 
-async def test_unload_entry(hass: HomeAssistant, entry, mock_create_batch) -> None:
+async def test_unload_entry(
+    hass: HomeAssistant, entry: MockConfigEntry, mock_create_batch: MagicMock
+) -> None:
     """Test being able to unload an entry.
 
     Queue should be empty, so adding events to the batch should not be called,
@@ -73,7 +75,7 @@ async def test_unload_entry(hass: HomeAssistant, entry, mock_create_batch) -> No
 
 
 async def test_failed_test_connection(
-    hass: HomeAssistant, mock_get_eventhub_properties
+    hass: HomeAssistant, mock_get_eventhub_properties: AsyncMock
 ) -> None:
     """Test being able to unload an entry."""
     entry = MockConfigEntry(
@@ -89,7 +91,9 @@ async def test_failed_test_connection(
 
 
 async def test_send_batch_error(
-    hass: HomeAssistant, entry_with_one_event, mock_send_batch
+    hass: HomeAssistant,
+    entry_with_one_event: MockConfigEntry,
+    mock_send_batch: AsyncMock,
 ) -> None:
     """Test a error in send_batch, including recovering at the next interval."""
     mock_send_batch.reset_mock()
@@ -111,7 +115,9 @@ async def test_send_batch_error(
 
 
 async def test_late_event(
-    hass: HomeAssistant, entry_with_one_event, mock_create_batch
+    hass: HomeAssistant,
+    entry_with_one_event: MockConfigEntry,
+    mock_create_batch: MagicMock,
 ) -> None:
     """Test the check on late events."""
     with patch(
@@ -128,7 +134,9 @@ async def test_late_event(
 
 
 async def test_full_batch(
-    hass: HomeAssistant, entry_with_one_event, mock_create_batch
+    hass: HomeAssistant,
+    entry_with_one_event: MockConfigEntry,
+    mock_create_batch: MagicMock,
 ) -> None:
     """Test the full batch behaviour."""
     mock_create_batch.add.side_effect = [ValueError, None]
@@ -208,7 +216,12 @@ async def test_full_batch(
     ],
     ids=["allowlist", "denylist", "filtered_allowlist", "filtered_denylist"],
 )
-async def test_filter(hass: HomeAssistant, entry, tests, mock_create_batch) -> None:
+async def test_filter(
+    hass: HomeAssistant,
+    entry: MockConfigEntry,
+    tests: list[FilterTest],
+    mock_create_batch: MagicMock,
+) -> None:
     """Test different filters.
 
     Filter_schema is also a fixture which is replaced by the filter_schema