diff --git a/roborock/data/v1/v1_code_mappings.py b/roborock/data/v1/v1_code_mappings.py index 2bc42481..c0bab14b 100644 --- a/roborock/data/v1/v1_code_mappings.py +++ b/roborock/data/v1/v1_code_mappings.py @@ -587,17 +587,6 @@ class RoborockDockDustCollectionModeCode(RoborockEnum): max = 4 -class RoborockDockWashTowelModeCode(RoborockEnum): - """Describes the wash towel mode of the vacuum cleaner.""" - - # TODO: Get the correct values for various different docks - unknown = -9999 - light = 0 - balanced = 1 - deep = 2 - smart = 10 - - class RoborockStateCode(RoborockEnum): unknown = 0 starting = 1 diff --git a/roborock/data/v1/v1_containers.py b/roborock/data/v1/v1_containers.py index 5285ace3..0d03a23b 100644 --- a/roborock/data/v1/v1_containers.py +++ b/roborock/data/v1/v1_containers.py @@ -39,6 +39,7 @@ from roborock.exceptions import RoborockException from ..containers import RoborockBase, RoborockBaseTimer, _attr_repr +from .v1_clean_modes import WashTowelModes from .v1_code_mappings import ( CleanFluidStatus, ClearWaterBoxStatus, @@ -48,7 +49,6 @@ RoborockDockDustCollectionModeCode, RoborockDockErrorCode, RoborockDockTypeCode, - RoborockDockWashTowelModeCode, RoborockErrorCode, RoborockFanPowerCode, RoborockFanSpeedP10, @@ -593,7 +593,7 @@ class DustCollectionMode(RoborockBase): @dataclass class WashTowelMode(RoborockBase): - wash_mode: RoborockDockWashTowelModeCode | None = None + wash_mode: WashTowelModes | None = None @dataclass diff --git a/roborock/devices/traits/v1/__init__.py b/roborock/devices/traits/v1/__init__.py index 7438c6b0..f6d121f1 100644 --- a/roborock/devices/traits/v1/__init__.py +++ b/roborock/devices/traits/v1/__init__.py @@ -188,14 +188,13 @@ def __init__( self._map_rpc_channel = map_rpc_channel self._web_api = web_api self._device_cache = device_cache - + self.device_features = DeviceFeaturesTrait(product, self._device_cache) self.status = StatusTrait(product) self.consumables = ConsumableTrait() self.rooms = RoomsTrait(home_data) self.maps = MapsTrait(self.status) self.map_content = MapContentTrait(map_parser_config) self.home = HomeTrait(self.status, self.maps, self.map_content, self.rooms, self._device_cache) - self.device_features = DeviceFeaturesTrait(product, self._device_cache) self.network_info = NetworkInfoTrait(device_uid, self._device_cache) self.routines = RoutinesTrait(device_uid, web_api) @@ -246,7 +245,10 @@ async def discover_features(self) -> None: _LOGGER.debug("Trait '%s' not supported, skipping", item.name) continue _LOGGER.debug("Trait '%s' is supported, initializing", item.name) - trait = item_type() + if item_type is WashTowelModeTrait: + trait = item_type(self.device_features) + else: + trait = item_type() setattr(self, item.name, trait) trait._rpc_channel = self._get_rpc_channel(trait) diff --git a/roborock/devices/traits/v1/wash_towel_mode.py b/roborock/devices/traits/v1/wash_towel_mode.py index 715f1ee9..19ae4ee0 100644 --- a/roborock/devices/traits/v1/wash_towel_mode.py +++ b/roborock/devices/traits/v1/wash_towel_mode.py @@ -1,8 +1,11 @@ """Trait for wash towel mode.""" -from roborock.data import WashTowelMode +from functools import cached_property + +from roborock.data import WashTowelMode, WashTowelModes, get_wash_towel_modes from roborock.device_features import is_wash_n_fill_dock from roborock.devices.traits.v1 import common +from roborock.devices.traits.v1.device_features import DeviceFeaturesTrait from roborock.roborock_typing import RoborockCommand @@ -11,3 +14,30 @@ class WashTowelModeTrait(WashTowelMode, common.V1TraitMixin): command = RoborockCommand.GET_WASH_TOWEL_MODE requires_dock_type = is_wash_n_fill_dock + + def __init__( + self, + device_feature_trait: DeviceFeaturesTrait | None = None, + wash_mode: WashTowelModes | None = None, + ) -> None: + super().__init__() + self.device_feature_trait = device_feature_trait + self.wash_mode = wash_mode + + @cached_property + def wash_towel_mode_options(self) -> list[WashTowelModes]: + if self.device_feature_trait is None: + return [] + return get_wash_towel_modes(self.device_feature_trait) + + async def set_wash_towel_mode(self, mode: WashTowelModes) -> None: + """Set the wash towel mode.""" + await self.rpc_channel.send_command(RoborockCommand.SET_WASH_TOWEL_MODE, params={"wash_mode": mode.code}) + + async def start_wash(self) -> None: + """Start washing the mop.""" + await self.rpc_channel.send_command(RoborockCommand.APP_START_WASH) + + async def stop_wash(self) -> None: + """Stop washing the mop.""" + await self.rpc_channel.send_command(RoborockCommand.APP_STOP_WASH) diff --git a/tests/devices/traits/v1/fixtures.py b/tests/devices/traits/v1/fixtures.py index 60fdfc42..8ad746ef 100644 --- a/tests/devices/traits/v1/fixtures.py +++ b/tests/devices/traits/v1/fixtures.py @@ -124,3 +124,4 @@ async def discover_features_fixture( await device.connect() assert device.v1_properties.status.dock_type == dock_type_code mock_rpc_channel.send_command.reset_mock() + mock_rpc_channel.send_command.side_effect = None diff --git a/tests/devices/traits/v1/test_wash_towel_mode.py b/tests/devices/traits/v1/test_wash_towel_mode.py index bda12124..8372a935 100644 --- a/tests/devices/traits/v1/test_wash_towel_mode.py +++ b/tests/devices/traits/v1/test_wash_towel_mode.py @@ -4,12 +4,15 @@ import pytest -from roborock.data import RoborockDockTypeCode, RoborockDockWashTowelModeCode +from roborock.data import ( + RoborockDockTypeCode, + WashTowelModes, +) from roborock.devices.device import RoborockDevice from roborock.devices.traits.v1.wash_towel_mode import WashTowelModeTrait from roborock.roborock_typing import RoborockCommand -WASH_TOWEL_MODE_DATA = [{"wash_mode": RoborockDockWashTowelModeCode.smart}] +WASH_TOWEL_MODE_DATA = {"wash_mode": WashTowelModes.SMART.code} @pytest.fixture(name="wash_towel_mode") @@ -50,7 +53,7 @@ async def test_wash_towel_mode_available( ] ) - assert wash_towel_mode.wash_mode == RoborockDockWashTowelModeCode.smart + assert wash_towel_mode.wash_mode == WashTowelModes.SMART @pytest.mark.parametrize( @@ -65,3 +68,133 @@ async def test_unsupported_wash_towel_mode( ) -> None: """Test that the trait is not available for unsupported dock types.""" assert wash_towel_mode is None + + +@pytest.mark.parametrize( + ("dock_type_code"), + [(RoborockDockTypeCode.s8_dock)], +) +@pytest.mark.parametrize( + ("wash_mode"), + [ + (WashTowelModes.SMART), + (WashTowelModes.LIGHT), + ], +) +async def test_set_wash_towel_mode( + wash_towel_mode: WashTowelModeTrait | None, + mock_rpc_channel: AsyncMock, + wash_mode: WashTowelModes, + dock_type_code: RoborockDockTypeCode, +) -> None: + """Test setting the wash towel mode.""" + assert wash_towel_mode is not None + + await wash_towel_mode.set_wash_towel_mode(wash_mode) + + mock_rpc_channel.send_command.assert_called_with( + RoborockCommand.SET_WASH_TOWEL_MODE, params={"wash_mode": wash_mode.code} + ) + + +@pytest.mark.parametrize( + ("dock_type_code"), + [(RoborockDockTypeCode.s8_dock)], +) +async def test_start_wash( + wash_towel_mode: WashTowelModeTrait | None, + mock_rpc_channel: AsyncMock, + dock_type_code: RoborockDockTypeCode, +) -> None: + """Test starting the wash.""" + assert wash_towel_mode is not None + + await wash_towel_mode.start_wash() + + mock_rpc_channel.send_command.assert_called_with(RoborockCommand.APP_START_WASH) + + +@pytest.mark.parametrize( + ("dock_type_code"), + [(RoborockDockTypeCode.s8_dock)], +) +async def test_stop_wash( + wash_towel_mode: WashTowelModeTrait | None, + mock_rpc_channel: AsyncMock, + dock_type_code: RoborockDockTypeCode, +) -> None: + """Test stopping the wash.""" + assert wash_towel_mode is not None + + await wash_towel_mode.stop_wash() + + mock_rpc_channel.send_command.assert_called_with(RoborockCommand.APP_STOP_WASH) + + +@pytest.mark.parametrize( + ("dock_type_code"), + [(RoborockDockTypeCode.s8_dock)], +) +@pytest.mark.parametrize( + ( + "is_super_deep_wash_supported", + "is_dirty_replenish_clean_supported", + "expected_modes", + ), + [ + ( + False, + False, + [WashTowelModes.LIGHT, WashTowelModes.BALANCED, WashTowelModes.DEEP], + ), + ( + True, + False, + [ + WashTowelModes.LIGHT, + WashTowelModes.BALANCED, + WashTowelModes.DEEP, + WashTowelModes.SUPER_DEEP, + ], + ), + ( + False, + True, + [ + WashTowelModes.LIGHT, + WashTowelModes.BALANCED, + WashTowelModes.DEEP, + WashTowelModes.SMART, + ], + ), + ( + True, + True, + [ + WashTowelModes.LIGHT, + WashTowelModes.BALANCED, + WashTowelModes.DEEP, + WashTowelModes.SMART, + ], + ), + ], +) +async def test_wash_towel_mode_options( + wash_towel_mode: WashTowelModeTrait | None, + dock_type_code: RoborockDockTypeCode, + is_super_deep_wash_supported: bool, + is_dirty_replenish_clean_supported: bool, + expected_modes: list[WashTowelModes], +) -> None: + """Test what modes are available based on device features.""" + assert wash_towel_mode is not None + # We need to clear the cached property to ensure it re-reads the features + if "wash_towel_mode_options" in wash_towel_mode.__dict__: + del wash_towel_mode.__dict__["wash_towel_mode_options"] + + # Mock the device features + assert wash_towel_mode.device_feature_trait is not None + wash_towel_mode.device_feature_trait.is_super_deep_wash_supported = is_super_deep_wash_supported + wash_towel_mode.device_feature_trait.is_dirty_replenish_clean_supported = is_dirty_replenish_clean_supported + + assert wash_towel_mode.wash_towel_mode_options == expected_modes