fix: Remove DockSummary and make dust collection mode optional based on dock type

Author: allenporter

roborock/device_features.py

@@ -642,3 +642,13 @@ def get_supported_features(self) -> list[str]:
     RoborockDockTypeCode.saros_r10_dock,
     RoborockDockTypeCode.qrevo_curv_dock,
 ]
+
+
+def is_wash_n_fill_dock(dock_type: RoborockDockTypeCode) -> bool:
+    """Check if the dock type is a wash and fill dock."""
+    return dock_type in WASH_N_FILL_DOCK_TYPES
+
+
+def is_valid_dock(dock_type: RoborockDockTypeCode) -> bool:
+    """Check if device supports a dock."""
+    return dock_type != RoborockDockTypeCode.no_dock

roborock/devices/traits/v1/__init__.py

@@ -26,9 +26,8 @@
     - `requires_feature` - The string name of the device feature that must be supported
         for this trait to be enabled. See `DeviceFeaturesTrait` for a list of
         available features.
-    - `requires_dock_type` - If set, this trait requires the
-        device to have any of the specified dock types to be enabled. See
-        `RoborockDockTypeCode` for a list of available dock types.
+    - `requires_dock_type` - If set, this is a function that accepts a `RoborockDockTypeCode`
+        and returns a boolean indicating whether the trait is supported for that dock type.
 """
 
 import logging
@@ -50,7 +49,6 @@
 from .consumeable import ConsumableTrait
 from .device_features import DeviceFeaturesTrait
 from .do_not_disturb import DoNotDisturbTrait
-from .dock_summary import DockSummaryTrait
 from .dust_collection_mode import DustCollectionModeTrait
 from .flow_led_status import FlowLedStatusTrait
 from .home import HomeTrait
@@ -84,7 +82,6 @@
     "FlowLedStatusTrait",
     "LedStatusTrait",
     "ValleyElectricityTimerTrait",
-    "DockSummaryTrait",
     "DustCollectionModeTrait",
     "WashTowelModeTrait",
     "SmartWashParamsTrait",
@@ -111,16 +108,15 @@ class PropertiesApi(Trait):
     consumables: ConsumableTrait
     home: HomeTrait
     device_features: DeviceFeaturesTrait
-    dust_collection_mode: DustCollectionModeTrait
 
     # Optional features that may not be supported on all devices
     child_lock: ChildLockTrait | None = None
     led_status: LedStatusTrait | None = None
     flow_led_status: FlowLedStatusTrait | None = None
     valley_electricity_timer: ValleyElectricityTimerTrait | None = None
+    dust_collection_mode: DustCollectionModeTrait | None = None
     wash_towel_mode: WashTowelModeTrait | None = None
     smart_wash_params: SmartWashParamsTrait | None = None
-    dock_summary: DockSummaryTrait | None = None
 
     def __init__(
         self,
@@ -191,35 +187,21 @@ async def discover_features(self) -> None:
 
             # Union args may not be in declared order
             item_type = union_args[0] if union_args[1] is type(None) else union_args[1]
-            if item_type is DockSummaryTrait:
-                # DockSummaryTrait is created manually below
-                continue
-            trait = item_type()
-            if not self._is_supported(trait, item.name, dock_type):
+            if not self._is_supported(item_type, item.name, dock_type):
                 _LOGGER.debug("Trait '%s' not supported, skipping", item.name)
                 continue
-
             _LOGGER.debug("Trait '%s' is supported, initializing", item.name)
+            trait = item_type()
             setattr(self, item.name, trait)
             trait._rpc_channel = self._get_rpc_channel(trait)
 
-        if dock_type is None or dock_type == RoborockDockTypeCode.no_dock:
-            _LOGGER.debug("Trait 'dock_summary' not supported, skipping")
-        else:
-            _LOGGER.debug("Trait 'dock_summary' is supported, initializing")
-            self.dock_summary = DockSummaryTrait(
-                self.dust_collection_mode,
-                self.wash_towel_mode,
-                self.smart_wash_params,
-            )
-
-    def _is_supported(self, trait: V1TraitMixin, name: str, dock_type: RoborockDockTypeCode) -> bool:
+    def _is_supported(self, trait_type: type[V1TraitMixin], name: str, dock_type: RoborockDockTypeCode) -> bool:
         """Check if a trait is supported by the device."""
 
-        if (requires_dock_type := getattr(trait, "requires_dock_type", None)) is not None:
-            return dock_type in requires_dock_type
+        if (requires_dock_type := getattr(trait_type, "requires_dock_type", None)) is not None:
+            return requires_dock_type(dock_type)
 
-        if (feature_name := getattr(trait, "requires_feature", None)) is None:
+        if (feature_name := getattr(trait_type, "requires_feature", None)) is None:
             _LOGGER.debug("Optional trait missing 'requires_feature' attribute %s, skipping", name)
             return False
         if (is_supported := getattr(self.device_features, feature_name)) is None:

roborock/devices/traits/v1/common.py

@@ -82,6 +82,7 @@ async def refresh(self) -> Self:
         new_data = self._parse_response(response)
         if not isinstance(new_data, RoborockBase):
             raise ValueError(f"Internal error, unexpected response type: {new_data!r}")
+        _LOGGER.debug("Refreshed %s: %s", self.__class__.__name__, new_data)
         self._update_trait_values(new_data)
         return self
 

roborock/devices/traits/v1/dock_summary.py

@@ -1,6 +1,7 @@
 """Trait for dust collection mode."""
 
 from roborock.containers import DustCollectionMode
+from roborock.device_features import is_valid_dock
 from roborock.devices.traits.v1 import common
 from roborock.roborock_typing import RoborockCommand
 
@@ -9,3 +10,4 @@ class DustCollectionModeTrait(DustCollectionMode, common.V1TraitMixin):
     """Trait for dust collection mode."""
 
     command = RoborockCommand.GET_DUST_COLLECTION_MODE
+    requires_dock_type = is_valid_dock

roborock/devices/traits/v1/dust_collection_mode.py

@@ -1,7 +1,7 @@
 """Trait for smart wash parameters."""
 
 from roborock.containers import SmartWashParams
-from roborock.device_features import WASH_N_FILL_DOCK_TYPES
+from roborock.device_features import is_wash_n_fill_dock
 from roborock.devices.traits.v1 import common
 from roborock.roborock_typing import RoborockCommand
 
@@ -10,4 +10,4 @@ class SmartWashParamsTrait(SmartWashParams, common.V1TraitMixin):
     """Trait for smart wash parameters."""
 
     command = RoborockCommand.GET_SMART_WASH_PARAMS
-    requires_dock_type = WASH_N_FILL_DOCK_TYPES
+    requires_dock_type = is_wash_n_fill_dock

roborock/devices/traits/v1/smart_wash_params.py

@@ -1,7 +1,7 @@
 """Trait for wash towel mode."""
 
 from roborock.containers import WashTowelMode
-from roborock.device_features import WASH_N_FILL_DOCK_TYPES
+from roborock.device_features import is_wash_n_fill_dock
 from roborock.devices.traits.v1 import common
 from roborock.roborock_typing import RoborockCommand
 
@@ -10,4 +10,4 @@ class WashTowelModeTrait(WashTowelMode, common.V1TraitMixin):
     """Trait for wash towel mode."""
 
     command = RoborockCommand.GET_WASH_TOWEL_MODE
-    requires_dock_type = WASH_N_FILL_DOCK_TYPES
+    requires_dock_type = is_wash_n_fill_dock

roborock/devices/traits/v1/wash_towel_mode.py

@@ -4,6 +4,7 @@
 
 import pytest
 
+from roborock.code_mappings import RoborockDockTypeCode
 from roborock.containers import HomeData, S7MaxVStatus, UserData
 from roborock.devices.cache import Cache, InMemoryCache
 from roborock.devices.device import RoborockDevice
@@ -68,3 +69,33 @@ def device_fixture(
             roborock_cache,
         ),
     )
+
+
+@pytest.fixture(name="dock_type_code", autouse=True)
+def dock_type_code_fixture(request: pytest.FixtureRequest) -> RoborockDockTypeCode | None:
+    """Fixture to provide the dock type code for parameterized tests."""
+    return RoborockDockTypeCode.s7_max_ultra_dock
+
+
+@pytest.fixture(autouse=True)
+async def discover_features_fixture(
+    device: RoborockDevice,
+    mock_rpc_channel: AsyncMock,
+    dock_type_code: RoborockDockTypeCode | None,
+) -> None:
+    """Fixture to set up the clean summary for tests.
+
+    The CleanRecordTrait depends on the CleanSummaryTrait, so we need to
+    prepare that first.
+    """
+    assert device.v1_properties
+    mock_rpc_channel.send_command.side_effect = [
+        [mock_data.APP_GET_INIT_STATUS],
+        {
+            **mock_data.STATUS,
+            "dock_type": dock_type_code,
+        },
+    ]
+    await device.v1_properties.discover_features()
+    assert device.v1_properties.status.dock_type == dock_type_code
+    mock_rpc_channel.send_command.reset_mock()