Add MockOf typing helper

tests/components/govee_light_local/conftest.py

@@ -2,18 +2,20 @@
 
 from asyncio import Event
 from collections.abc import Generator
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
 from govee_local_api import GoveeLightCapability
 import pytest
 
 from homeassistant.components.govee_light_local.coordinator import GoveeController
 
+from tests.typing import MockOf
+
 
 @pytest.fixture(name="mock_govee_api")
-def fixture_mock_govee_api():
+def fixture_mock_govee_api() -> MockOf[GoveeController]:
     """Set up Govee Local API fixture."""
-    mock_api = AsyncMock(spec=GoveeController)
+    mock_api = Mock(spec=GoveeController)
     mock_api.start = AsyncMock()
     mock_api.cleanup = MagicMock(return_value=Event())
     mock_api.cleanup.return_value.set()

tests/components/govee_light_local/test_config_flow.py

@@ -3,7 +3,7 @@
 from errno import EADDRINUSE
 from unittest.mock import AsyncMock, patch
 
-from govee_local_api import GoveeDevice
+from govee_local_api import GoveeController, GoveeDevice
 
 from homeassistant import config_entries
 from homeassistant.components.govee_light_local.const import DOMAIN
@@ -12,8 +12,10 @@
 
 from .conftest import DEFAULT_CAPABILITEIS
 
+from tests.typing import MockOf
 
-def _get_devices(mock_govee_api: AsyncMock) -> list[GoveeDevice]:
+
+def _get_devices(mock_govee_api: MockOf[GoveeController]) -> list[GoveeDevice]:
     return [
         GoveeDevice(
             controller=mock_govee_api,
@@ -26,7 +28,9 @@ def _get_devices(mock_govee_api: AsyncMock) -> list[GoveeDevice]:
 
 
 async def test_creating_entry_has_no_devices(
-    hass: HomeAssistant, mock_setup_entry: AsyncMock, mock_govee_api: AsyncMock
+    hass: HomeAssistant,
+    mock_setup_entry: AsyncMock,
+    mock_govee_api: MockOf[GoveeController],
 ) -> None:
     """Test setting up Govee with no devices."""
 
@@ -61,7 +65,7 @@ async def test_creating_entry_has_no_devices(
 async def test_creating_entry_has_with_devices(
     hass: HomeAssistant,
     mock_setup_entry: AsyncMock,
-    mock_govee_api: AsyncMock,
+    mock_govee_api: MockOf[GoveeController],
 ) -> None:
     """Test setting up Govee with devices."""
 
@@ -90,7 +94,7 @@ async def test_creating_entry_has_with_devices(
 async def test_creating_entry_errno(
     hass: HomeAssistant,
     mock_setup_entry: AsyncMock,
-    mock_govee_api: AsyncMock,
+    mock_govee_api: MockOf[GoveeController],
 ) -> None:
     """Test setting up Govee with devices."""
 

tests/components/govee_light_local/test_light.py

@@ -1,9 +1,9 @@
 """Test Govee light local."""
 
 from errno import EADDRINUSE, ENETDOWN
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import patch
 
-from govee_local_api import GoveeDevice
+from govee_local_api import GoveeController, GoveeDevice
 
 from homeassistant.components.govee_light_local.const import DOMAIN
 from homeassistant.components.light import ATTR_SUPPORTED_COLOR_MODES, ColorMode
@@ -13,10 +13,11 @@
 from .conftest import DEFAULT_CAPABILITEIS
 
 from tests.common import MockConfigEntry
+from tests.typing import MockOf
 
 
 async def test_light_known_device(
-    hass: HomeAssistant, mock_govee_api: AsyncMock
+    hass: HomeAssistant, mock_govee_api: MockOf[GoveeController]
 ) -> None:
     """Test adding a known device."""
 
@@ -55,7 +56,7 @@ async def test_light_known_device(
 
 
 async def test_light_unknown_device(
-    hass: HomeAssistant, mock_govee_api: AsyncMock
+    hass: HomeAssistant, mock_govee_api: MockOf[GoveeController]
 ) -> None:
     """Test adding an unknown device."""
 
@@ -87,7 +88,9 @@ async def test_light_unknown_device(
         assert light.attributes[ATTR_SUPPORTED_COLOR_MODES] == [ColorMode.ONOFF]
 
 
-async def test_light_remove(hass: HomeAssistant, mock_govee_api: AsyncMock) -> None:
+async def test_light_remove(
+    hass: HomeAssistant, mock_govee_api: MockOf[GoveeController]
+) -> None:
     """Test adding a known device."""
 
     mock_govee_api.devices = [
@@ -118,7 +121,7 @@ async def test_light_remove(hass: HomeAssistant, mock_govee_api: AsyncMock) -> N
 
 
 async def test_light_setup_retry(
-    hass: HomeAssistant, mock_govee_api: AsyncMock
+    hass: HomeAssistant, mock_govee_api: MockOf[GoveeController]
 ) -> None:
     """Test adding an unknown device."""
 
@@ -140,7 +143,7 @@ async def test_light_setup_retry(
 
 
 async def test_light_setup_retry_eaddrinuse(
-    hass: HomeAssistant, mock_govee_api: AsyncMock
+    hass: HomeAssistant, mock_govee_api: MockOf[GoveeController]
 ) -> None:
     """Test adding an unknown device."""
 
@@ -168,7 +171,7 @@ async def test_light_setup_retry_eaddrinuse(
 
 
 async def test_light_setup_error(
-    hass: HomeAssistant, mock_govee_api: AsyncMock
+    hass: HomeAssistant, mock_govee_api: MockOf[GoveeController]
 ) -> None:
     """Test adding an unknown device."""
 
@@ -195,7 +198,9 @@ async def test_light_setup_error(
         assert entry.state is ConfigEntryState.SETUP_ERROR
 
 
-async def test_light_on_off(hass: HomeAssistant, mock_govee_api: MagicMock) -> None:
+async def test_light_on_off(
+    hass: HomeAssistant, mock_govee_api: MockOf[GoveeController]
+) -> None:
     """Test adding a known device."""
 
     mock_govee_api.devices = [
@@ -252,7 +257,9 @@ async def test_light_on_off(hass: HomeAssistant, mock_govee_api: MagicMock) -> N
         mock_govee_api.turn_on_off.assert_awaited_with(mock_govee_api.devices[0], False)
 
 
-async def test_light_brightness(hass: HomeAssistant, mock_govee_api: MagicMock) -> None:
+async def test_light_brightness(
+    hass: HomeAssistant, mock_govee_api: MockOf[GoveeController]
+) -> None:
     """Test changing brightness."""
     mock_govee_api.devices = [
         GoveeDevice(
@@ -327,7 +334,9 @@ async def test_light_brightness(hass: HomeAssistant, mock_govee_api: MagicMock)
         )
 
 
-async def test_light_color(hass: HomeAssistant, mock_govee_api: MagicMock) -> None:
+async def test_light_color(
+    hass: HomeAssistant, mock_govee_api: MockOf[GoveeController]
+) -> None:
     """Test changing brightness."""
     mock_govee_api.devices = [
         GoveeDevice(

tests/typing.py

@@ -3,8 +3,8 @@
 from __future__ import annotations
 
 from collections.abc import Callable, Coroutine
-from typing import TYPE_CHECKING, Any
-from unittest.mock import MagicMock
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
+from unittest.mock import MagicMock, Mock
 
 from aiohttp import ClientWebSocketResponse
 from aiohttp.test_utils import TestClient
@@ -14,6 +14,12 @@
     # testcase which does not use the recorder.
     from homeassistant.components.recorder import Recorder
 
+_T = TypeVar("_T")
+
+
+class MockOf(Mock, Generic[_T]):
+    """Add ability to set underlying type for Mock."""
+
 
 class MockHAClientWebSocket(ClientWebSocketResponse):
     """Protocol for a wrapped ClientWebSocketResponse."""