Improve contextmanager typing (#122250)

homeassistant/components/zha/radio_manager.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import asyncio
+from collections.abc import AsyncIterator
 import contextlib
 from contextlib import suppress
 import copy
@@ -157,7 +158,7 @@ class ZhaRadioManager:
         return mgr
 
     @contextlib.asynccontextmanager
-    async def connect_zigpy_app(self) -> ControllerApplication:
+    async def connect_zigpy_app(self) -> AsyncIterator[ControllerApplication]:
         """Connect to the radio with the current config and then clean up."""
         assert self.radio_type is not None
 

tests/components/automation/test_blueprint.py

@@ -1,8 +1,10 @@
 """Test built-in blueprints."""
 
 import asyncio
+from collections.abc import Iterator
 import contextlib
 from datetime import timedelta
+from os import PathLike
 import pathlib
 from typing import Any
 from unittest.mock import patch
@@ -23,7 +25,9 @@ BUILTIN_BLUEPRINT_FOLDER = pathlib.Path(automation.__file__).parent / "blueprint
 
 
 @contextlib.contextmanager
-def patch_blueprint(blueprint_path: str, data_path):
+def patch_blueprint(
+    blueprint_path: str, data_path: str | PathLike[str]
+) -> Iterator[None]:
     """Patch blueprint loading from a different source."""
     orig_load = models.DomainBlueprints._load_blueprint
 

tests/components/bluetooth/test_wrappers.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from collections.abc import Iterator
 from contextlib import contextmanager
 from unittest.mock import patch
 
@@ -27,7 +28,7 @@ from . import _get_manager, generate_advertisement_data, generate_ble_device
 
 
 @contextmanager
-def mock_shutdown(manager: HomeAssistantBluetoothManager) -> None:
+def mock_shutdown(manager: HomeAssistantBluetoothManager) -> Iterator[None]:
     """Mock shutdown of the HomeAssistantBluetoothManager."""
     manager.shutdown = True
     yield

tests/components/homeassistant_sky_connect/test_config_flow.py

@@ -1,7 +1,7 @@
 """Test the Home Assistant SkyConnect config flow."""
 
 import asyncio
-from collections.abc import Awaitable, Callable
+from collections.abc import Awaitable, Callable, Iterator
 import contextlib
 from typing import Any
 from unittest.mock import AsyncMock, Mock, call, patch
@@ -80,7 +80,7 @@ def mock_addon_info(
         update_available=False,
         version=None,
     ),
-):
+) -> Iterator[tuple[Mock, Mock]]:
     """Mock the main addon states for the config flow."""
     mock_flasher_manager = Mock(spec_set=get_zigbee_flasher_addon_manager(hass))
     mock_flasher_manager.addon_name = "Silicon Labs Flasher"

tests/components/netatmo/common.py

@@ -1,9 +1,10 @@
 """Common methods used across tests for Netatmo."""
 
+from collections.abc import Iterator
 from contextlib import contextmanager
 import json
 from typing import Any
-from unittest.mock import AsyncMock, patch
+from unittest.mock import patch
 
 from syrupy import SnapshotAssertion
 
@@ -109,7 +110,7 @@ async def simulate_webhook(hass: HomeAssistant, webhook_id: str, response) -> No
 
 
 @contextmanager
-def selected_platforms(platforms: list[Platform]) -> AsyncMock:
+def selected_platforms(platforms: list[Platform]) -> Iterator[None]:
     """Restrict loaded platforms to list given."""
     with (
         patch("homeassistant.components.netatmo.data_handler.PLATFORMS", platforms),

tests/components/renault/conftest.py

@@ -1,6 +1,6 @@
 """Provide common Renault fixtures."""
 
-from collections.abc import Generator
+from collections.abc import Generator, Iterator
 import contextlib
 from types import MappingProxyType
 from typing import Any
@@ -200,7 +200,7 @@ def patch_fixtures_with_no_data():
 
 
 @contextlib.contextmanager
-def _patch_fixtures_with_side_effect(side_effect: Any):
+def _patch_fixtures_with_side_effect(side_effect: Any) -> Iterator[None]:
     """Mock fixtures."""
     with (
         patch(

tests/test_util/aiohttp.py

@@ -1,6 +1,7 @@
 """Aiohttp test utils."""
 
 import asyncio
+from collections.abc import Iterator
 from contextlib import contextmanager
 from http import HTTPStatus
 import re
@@ -296,7 +297,7 @@ class AiohttpClientMockResponse:
 
 
 @contextmanager
-def mock_aiohttp_client():
+def mock_aiohttp_client() -> Iterator[AiohttpClientMocker]:
     """Context manager to mock aiohttp client."""
     mocker = AiohttpClientMocker()