Update roborock test typing (home-assistant#157370)

Author: allenporter

tests/components/roborock/conftest.py

@@ -1,7 +1,7 @@
 """Global fixtures for Roborock integration."""
 
 import asyncio
-from collections.abc import Awaitable, Callable, Generator
+from collections.abc import Generator
 from copy import deepcopy
 import logging
 import pathlib
@@ -13,18 +13,34 @@
 from roborock import RoborockCategory
 from roborock.data import (
     CombinedMapInfo,
+    DnDTimer,
     DyadError,
     HomeDataDevice,
     HomeDataProduct,
     NamedRoomMapping,
     NetworkInfo,
+    RoborockBase,
     RoborockDyadStateCode,
     ZeoError,
     ZeoState,
 )
 from roborock.devices.device import RoborockDevice
-from roborock.devices.traits.v1.map_content import MapContent
-from roborock.devices.traits.v1.volume import SoundVolume
+from roborock.devices.traits.v1 import PropertiesApi
+from roborock.devices.traits.v1.clean_summary import CleanSummaryTrait
+from roborock.devices.traits.v1.command import CommandTrait
+from roborock.devices.traits.v1.common import V1TraitMixin
+from roborock.devices.traits.v1.consumeable import ConsumableTrait
+from roborock.devices.traits.v1.do_not_disturb import DoNotDisturbTrait
+from roborock.devices.traits.v1.dust_collection_mode import DustCollectionModeTrait
+from roborock.devices.traits.v1.home import HomeTrait
+from roborock.devices.traits.v1.map_content import MapContent, MapContentTrait
+from roborock.devices.traits.v1.maps import MapsTrait
+from roborock.devices.traits.v1.network_info import NetworkInfoTrait
+from roborock.devices.traits.v1.routines import RoutinesTrait
+from roborock.devices.traits.v1.smart_wash_params import SmartWashParamsTrait
+from roborock.devices.traits.v1.status import StatusTrait
+from roborock.devices.traits.v1.volume import SoundVolumeTrait
+from roborock.devices.traits.v1.wash_towel_mode import WashTowelModeTrait
 from roborock.roborock_message import RoborockDyadDataProtocol, RoborockZeoProtocol
 
 from homeassistant.components.roborock.const import (
@@ -130,71 +146,100 @@ async def get_devices(self) -> list[RoborockDevice]:
         return self._devices
 
 
-def make_fake_switch(obj: Any) -> Any:
-    """Update the fake object to emulate the switch trait behavior."""
-    obj.is_on = True
-    obj.enable = AsyncMock()
-    obj.enable.side_effect = lambda: setattr(obj, "is_on", True)
-    obj.disable = AsyncMock()
-    obj.disable.side_effect = lambda: setattr(obj, "is_on", False)
-    obj.refresh = AsyncMock()
-    return obj
+def make_mock_trait(
+    trait_spec: type[V1TraitMixin] | None = None,
+    dataclass_template: RoborockBase | None = None,
+) -> AsyncMock:
+    """Create a mock roborock trait."""
+    trait = AsyncMock(spec=trait_spec or V1TraitMixin)
+    if dataclass_template is not None:
+        # Copy all attributes and property methods (e.g. computed properties)
+        template_copy = deepcopy(dataclass_template)
+        for attr_name in dir(template_copy):
+            if attr_name.startswith("_"):
+                continue
+            setattr(trait, attr_name, getattr(template_copy, attr_name))
+    trait.refresh = AsyncMock()
+    return trait
+
+
+def make_mock_switch(
+    trait_spec: type[V1TraitMixin] | None = None,
+    dataclass_template: RoborockBase | None = None,
+) -> AsyncMock:
+    """Create a mock roborock switch trait."""
+    trait = make_mock_trait(
+        trait_spec=trait_spec,
+        dataclass_template=dataclass_template,
+    )
+    trait.is_on = True
+    trait.enable = AsyncMock()
+    trait.enable.side_effect = lambda: setattr(trait, "is_on", True)
+    trait.disable = AsyncMock()
+    trait.disable.side_effect = lambda: setattr(trait, "is_on", False)
+    return trait
 
 
-def set_timer_fn(obj: Any) -> Callable[[Any], Awaitable[None]]:
+def make_dnd_timer(dataclass_template: RoborockBase) -> AsyncMock:
     """Make a function for the fake timer trait that emulates the real behavior."""
+    dnd_trait = make_mock_switch(
+        trait_spec=DoNotDisturbTrait,
+        dataclass_template=dataclass_template,
+    )
 
-    async def update_timer_attributes(timer: Any) -> None:
-        setattr(obj, "start_hour", timer.start_hour)
-        setattr(obj, "start_minute", timer.start_minute)
-        setattr(obj, "end_hour", timer.end_hour)
-        setattr(obj, "end_minute", timer.end_minute)
-        setattr(obj, "enabled", timer.enabled)
+    async def set_dnd_timer(timer: DnDTimer) -> None:
+        setattr(dnd_trait, "start_hour", timer.start_hour)
+        setattr(dnd_trait, "start_minute", timer.start_minute)
+        setattr(dnd_trait, "end_hour", timer.end_hour)
+        setattr(dnd_trait, "end_minute", timer.end_minute)
+        setattr(dnd_trait, "enabled", timer.enabled)
 
-    return update_timer_attributes
+    dnd_trait.set_dnd_timer = AsyncMock()
+    dnd_trait.set_dnd_timer.side_effect = set_dnd_timer
+    return dnd_trait
 
 
-def create_v1_properties(network_info: NetworkInfo) -> Mock:
+def create_v1_properties(network_info: NetworkInfo) -> AsyncMock:
     """Create v1 properties for each fake device."""
-    v1_properties = Mock()
-    v1_properties.status: Any = deepcopy(STATUS)
-    v1_properties.status.refresh = AsyncMock()
-    v1_properties.dnd: Any = make_fake_switch(deepcopy(DND_TIMER))
-    v1_properties.dnd.set_dnd_timer = AsyncMock()
-    v1_properties.dnd.set_dnd_timer.side_effect = set_timer_fn(v1_properties.dnd)
-    v1_properties.clean_summary: Any = deepcopy(CLEAN_SUMMARY)
+    v1_properties = AsyncMock(spec=PropertiesApi)
+    v1_properties.status = make_mock_trait(
+        trait_spec=StatusTrait,
+        dataclass_template=STATUS,
+    )
+    v1_properties.dnd = make_dnd_timer(dataclass_template=DND_TIMER)
+    v1_properties.clean_summary = make_mock_trait(
+        trait_spec=CleanSummaryTrait,
+        dataclass_template=CLEAN_SUMMARY,
+    )
     v1_properties.clean_summary.last_clean_record = deepcopy(CLEAN_RECORD)
-    v1_properties.clean_summary.refresh = AsyncMock()
-    v1_properties.consumables = deepcopy(CONSUMABLE)
-    v1_properties.consumables.refresh = AsyncMock()
+    v1_properties.consumables = make_mock_trait(
+        trait_spec=ConsumableTrait, dataclass_template=CONSUMABLE
+    )
     v1_properties.consumables.reset_consumable = AsyncMock()
-    v1_properties.sound_volume = SoundVolume(volume=50)
+    v1_properties.sound_volume = make_mock_trait(trait_spec=SoundVolumeTrait)
+    v1_properties.sound_volume.volume = 50
     v1_properties.sound_volume.set_volume = AsyncMock()
     v1_properties.sound_volume.set_volume.side_effect = lambda vol: setattr(
         v1_properties.sound_volume, "volume", vol
     )
-    v1_properties.sound_volume.refresh = AsyncMock()
-    v1_properties.command = AsyncMock()
+    v1_properties.command = AsyncMock(spec=CommandTrait)
     v1_properties.command.send = AsyncMock()
-    v1_properties.maps = AsyncMock()
+    v1_properties.maps = make_mock_trait(trait_spec=MapsTrait)
     v1_properties.maps.current_map = MULTI_MAP_LIST.map_info[1].map_flag
-    v1_properties.maps.refresh = AsyncMock()
     v1_properties.maps.set_current_map = AsyncMock()
-    v1_properties.map_content = AsyncMock()
+    v1_properties.map_content = make_mock_trait(trait_spec=MapContentTrait)
     v1_properties.map_content.image_content = b"\x89PNG-001"
     v1_properties.map_content.map_data = deepcopy(MAP_DATA)
-    v1_properties.map_content.refresh = AsyncMock()
-    v1_properties.child_lock = make_fake_switch(AsyncMock())
-    v1_properties.led_status = make_fake_switch(AsyncMock())
-    v1_properties.flow_led_status = make_fake_switch(AsyncMock())
-    v1_properties.valley_electricity_timer = make_fake_switch(AsyncMock())
-    v1_properties.dust_collection_mode = AsyncMock()
-    v1_properties.dust_collection_mode.refresh = AsyncMock()
-    v1_properties.wash_towel_mode = AsyncMock()
-    v1_properties.wash_towel_mode.refresh = AsyncMock()
-    v1_properties.smart_wash_params = AsyncMock()
-    v1_properties.smart_wash_params.refresh = AsyncMock()
-    v1_properties.home = AsyncMock()
+    v1_properties.child_lock = make_mock_switch()
+    v1_properties.led_status = make_mock_switch()
+    v1_properties.flow_led_status = make_mock_switch()
+    v1_properties.valley_electricity_timer = make_mock_switch()
+    v1_properties.dust_collection_mode = make_mock_trait(
+        trait_spec=DustCollectionModeTrait
+    )
+    v1_properties.wash_towel_mode = make_mock_trait(trait_spec=WashTowelModeTrait)
+    v1_properties.smart_wash_params = make_mock_trait(trait_spec=SmartWashParamsTrait)
+    v1_properties.home = make_mock_trait(trait_spec=HomeTrait)
     home_map_info = {
         map_data.map_flag: CombinedMapInfo(
             name=map_data.name,
@@ -219,10 +264,11 @@ def create_v1_properties(network_info: NetworkInfo) -> Mock:
     v1_properties.home.home_map_info = home_map_info
     v1_properties.home.current_map_data = home_map_info[STATUS.current_map]
     v1_properties.home.home_map_content = home_map_content
-    v1_properties.home.refresh = AsyncMock()
-    v1_properties.network_info = deepcopy(network_info)
-    v1_properties.network_info.refresh = AsyncMock()
-    v1_properties.routines = AsyncMock()
+    v1_properties.network_info = make_mock_trait(
+        trait_spec=NetworkInfoTrait,
+        dataclass_template=network_info,
+    )
+    v1_properties.routines = make_mock_trait(trait_spec=RoutinesTrait)
     v1_properties.routines.get_routines = AsyncMock(return_value=SCENES)
     v1_properties.routines.execute_routine = AsyncMock()
     # Mock diagnostics for a subset of properties
@@ -267,21 +313,22 @@ def fake_vacuum_fixture(fake_devices: list[FakeDevice]) -> FakeDevice:
     return fake_devices[0]
 
 
-@pytest.fixture(name="send_message_side_effect")
-def send_message_side_effect_fixture() -> Any:
+@pytest.fixture(name="send_message_exception")
+def send_message_exception_fixture() -> Exception | None:
     """Fixture to return a side effect for the send_message method."""
     return None
 
 
 @pytest.fixture(name="vacuum_command", autouse=True)
 def fake_vacuum_command_fixture(
-    fake_vacuum: FakeDevice, send_message_side_effect: Any
-) -> Mock:
+    fake_vacuum: FakeDevice,
+    send_message_exception: Exception | None,
+) -> AsyncMock:
     """Get the fake vacuum device command trait for asserting that commands happened."""
     assert fake_vacuum.v1_properties is not None
     command_trait = fake_vacuum.v1_properties.command
-    if send_message_side_effect is not None:
-        command_trait.send.side_effect = send_message_side_effect
+    if send_message_exception is not None:
+        command_trait.send.side_effect = send_message_exception
     return command_trait
 
 

tests/components/roborock/test_switch.py

@@ -86,7 +86,7 @@ async def test_update_success(
     ],
 )
 @pytest.mark.parametrize(
-    "send_message_side_effect", [roborock.exceptions.RoborockTimeout]
+    "send_message_exception", [roborock.exceptions.RoborockTimeout]
 )
 async def test_update_failed(
     hass: HomeAssistant,