feat: add stronger mop washing support

Author: Lash-L

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

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

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)
 

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,29 @@ 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:
+        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)

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

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