feat: add stronger mop washing support (#765)

* feat: add stronger mop washing support

* chore: call super init

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* chore: simplify device_feature_trait usage

* chore: address test comments

* chore: move outside of loop

* chore: clean up and better match status trait

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

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 NamedRoomMapping, 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,
@@ -750,7 +750,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

@@ -234,6 +234,12 @@ async def discover_features(self) -> None:
         # Dock type also acts like a device feature for some traits.
         dock_type = await self._dock_type()
 
+        # Initialize traits with special arguments before the generic loop
+        if self.wash_towel_mode is None and self._is_supported(WashTowelModeTrait, "wash_towel_mode", dock_type):
+            wash_towel_mode = WashTowelModeTrait(self.device_features)
+            wash_towel_mode._rpc_channel = self._get_rpc_channel(wash_towel_mode)  # type: ignore[assignment]
+            self.wash_towel_mode = wash_towel_mode
+
         # Dynamically create any traits that need to be populated
         for item in fields(self):
             if (trait := getattr(self, item.name, None)) is not None:

roborock/devices/traits/v1/wash_towel_mode.py

@@ -1,8 +1,12 @@
 """Trait for wash towel mode."""
 
-from roborock.data import WashTowelMode
+from functools import cached_property
+from typing import Self
+
+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 +15,34 @@ 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:
+        super().__init__()
+        self.device_feature_trait = device_feature_trait
+
+    def _parse_response(self, response: common.V1ResponseData) -> Self:
+        """Parse the response from the device into a WashTowelMode object."""
+        if isinstance(response, list):
+            response = response[0]
+        if isinstance(response, dict):
+            return WashTowelMode.from_dict(response)
+        raise ValueError(f"Unexpected wash towel mode format: {response!r}")
+
+    @cached_property
+    def wash_towel_mode_options(self) -> list[WashTowelModes]:
+        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

@@ -126,3 +126,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")
@@ -31,7 +34,7 @@ def wash_towel_mode_trait(
     ],
 )
 async def test_wash_towel_mode_available(
-    wash_towel_mode: WashTowelModeTrait | None,
+    wash_towel_mode: WashTowelModeTrait,
     mock_rpc_channel: AsyncMock,
     dock_type_code: RoborockDockTypeCode,
 ) -> None:
@@ -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,130 @@ 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,
+    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,
+    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,
+    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,
+    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
+
+    # 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