feat: Add clean record and dock related traits (#550)

* feat: Add dock summary and clean record traits

This creates a new dock type required feature. This requires the status is fetched first initially.
The dock type is added to the cache since it is not expected to change once discovered.

The clean record trait depends on the clean summary and fetches the last clean record
from the summary.

Tests are added for clean record and dock summary. The dock summary tests also exercise
the feature discovery and underlying modes.

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

Author: allenporter

roborock/cli.py

@@ -119,6 +119,7 @@ class ConnectionCache(RoborockBase):
     home_data: HomeData | None = None
     network_info: dict[str, NetworkInfo] | None = None
     home_cache: dict[int, CombinedMapInfo] | None = None
+    trait_data: dict[str, Any] | None = None
 
 
 class DeviceConnectionManager:
@@ -267,6 +268,7 @@ async def get(self) -> CacheData:
             home_data=connection_cache.home_data,
             network_info=connection_cache.network_info or {},
             home_cache=connection_cache.home_cache,
+            trait_data=connection_cache.trait_data or {},
         )
 
     async def set(self, value: CacheData) -> None:
@@ -276,6 +278,7 @@ async def set(self, value: CacheData) -> None:
         connection_cache.home_data = value.home_data
         connection_cache.network_info = value.network_info
         connection_cache.home_cache = value.home_cache
+        connection_cache.trait_data = value.trait_data
         self.update(connection_cache)
 
 
@@ -401,9 +404,11 @@ async def _v1_trait(context: RoborockContext, device_id: str, display_func: Call
     device_manager = await context.get_device_manager()
     device = await device_manager.get_device(device_id)
     if device.v1_properties is None:
-        raise RoborockException(f"Device {device.name} does not support V1 protocol")
+        raise RoborockUnsupportedFeature(f"Device {device.name} does not support V1 protocol")
     await device.v1_properties.discover_features()
     trait = display_func(device.v1_properties)
+    if trait is None:
+        raise RoborockUnsupportedFeature("Trait not supported by device")
     await trait.refresh()
     return trait
 
@@ -440,6 +445,26 @@ async def clean_summary(ctx, device_id: str):
     await _display_v1_trait(context, device_id, lambda v1: v1.clean_summary)
 
 
+@session.command()
+@click.option("--device_id", required=True)
+@click.pass_context
+@async_command
+async def clean_record(ctx, device_id: str):
+    """Get device last clean record."""
+    context: RoborockContext = ctx.obj
+    await _display_v1_trait(context, device_id, lambda v1: v1.clean_record)
+
+
+@session.command()
+@click.option("--device_id", required=True)
+@click.pass_context
+@async_command
+async def dock_summary(ctx, device_id: str):
+    """Get device dock summary."""
+    context: RoborockContext = ctx.obj
+    await _display_v1_trait(context, device_id, lambda v1: v1.dock_summary)
+
+
 @session.command()
 @click.option("--device_id", required=True)
 @click.pass_context
@@ -938,6 +963,8 @@ def write_markdown_table(product_features: dict[str, dict[str, any]], all_featur
 cli.add_command(get_device_info)
 cli.add_command(update_docs)
 cli.add_command(clean_summary)
+cli.add_command(clean_record)
+cli.add_command(dock_summary)
 cli.add_command(volume)
 cli.add_command(set_volume)
 cli.add_command(maps)

roborock/device_features.py

@@ -4,7 +4,7 @@
 from enum import IntEnum, StrEnum
 from typing import Any
 
-from .code_mappings import RoborockProductNickname
+from .code_mappings import RoborockDockTypeCode, RoborockProductNickname
 from .containers import RoborockBase
 
 
@@ -630,3 +630,25 @@ def from_feature_flags(
     def get_supported_features(self) -> list[str]:
         """Returns a list of supported features (Primarily used for logging purposes)."""
         return [k for k, v in vars(self).items() if v]
+
+
+WASH_N_FILL_DOCK_TYPES = [
+    RoborockDockTypeCode.empty_wash_fill_dock,
+    RoborockDockTypeCode.s8_dock,
+    RoborockDockTypeCode.p10_dock,
+    RoborockDockTypeCode.p10_pro_dock,
+    RoborockDockTypeCode.s8_maxv_ultra_dock,
+    RoborockDockTypeCode.qrevo_s_dock,
+    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/cache.py

@@ -6,7 +6,7 @@
 """
 
 from dataclasses import dataclass, field
-from typing import Protocol
+from typing import Any, Protocol
 
 from roborock.containers import CombinedMapInfo, HomeData, NetworkInfo
 from roborock.device_features import DeviceFeatures
@@ -28,6 +28,9 @@ class CacheData:
     device_features: DeviceFeatures | None = None
     """Device features information."""
 
+    trait_data: dict[str, Any] | None = None
+    """Trait-specific cached data used internally for caching device features."""
+
 
 class Cache(Protocol):
     """Protocol for a cache that can store and retrieve values."""

roborock/devices/traits/v1/__init__.py

@@ -1,31 +1,66 @@
-"""Create traits for V1 devices."""
+"""Create traits for V1 devices.
+
+Traits are modular components that encapsulate specific features of a Roborock
+device. This module provides a factory function to create and initialize the
+appropriate traits for V1 devices based on their capabilities. They can also
+be considered groups of commands and parsing logic for that command.
+
+Traits have a `refresh()` method that can be called to update their state
+from the device. Some traits may also provide additional methods for modifying
+the device state.
+
+The most common pattern for a trait is to subclass `V1TraitMixin` and a `RoborockBase`
+dataclass, and define a `command` class variable that specifies the `RoborockCommand`
+used to fetch the trait data from the device. See `common.py` for more details
+on common patterns used across traits.
+
+There are some additional decorators in `common.py` that can be used to specify which
+RPC channel to use for the trait (standard, MQTT/cloud, or map-specific).
+
+  - `@common.mqtt_rpc_channel` - Use the MQTT RPC channel for this trait.
+  - `@common.map_rpc_channel` - Use the map RPC channel for this trait.
+
+There are also some attributes that specify device feature dependencies for
+optional traits:
+
+    - `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 is a function that accepts a `RoborockDockTypeCode`
+        and returns a boolean indicating whether the trait is supported for that dock type.
+"""
 
 import logging
 from dataclasses import dataclass, field, fields
-from typing import get_args
+from typing import Any, get_args
 
+from roborock.code_mappings import RoborockDockTypeCode
 from roborock.containers import HomeData, HomeDataProduct
 from roborock.devices.cache import Cache
 from roborock.devices.traits import Trait
 from roborock.devices.v1_rpc_channel import V1RpcChannel
 from roborock.map.map_parser import MapParserConfig
 
 from .child_lock import ChildLockTrait
+from .clean_record import CleanRecordTrait
 from .clean_summary import CleanSummaryTrait
 from .command import CommandTrait
 from .common import V1TraitMixin
 from .consumeable import ConsumableTrait
 from .device_features import DeviceFeaturesTrait
 from .do_not_disturb import DoNotDisturbTrait
+from .dust_collection_mode import DustCollectionModeTrait
 from .flow_led_status import FlowLedStatusTrait
 from .home import HomeTrait
 from .led_status import LedStatusTrait
 from .map_content import MapContentTrait
 from .maps import MapsTrait
 from .rooms import RoomsTrait
+from .smart_wash_params import SmartWashParamsTrait
 from .status import StatusTrait
 from .valley_electricity_timer import ValleyElectricityTimerTrait
 from .volume import SoundVolumeTrait
+from .wash_towel_mode import WashTowelModeTrait
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -35,6 +70,7 @@
     "StatusTrait",
     "DoNotDisturbTrait",
     "CleanSummaryTrait",
+    "CleanRecordTrait",
     "SoundVolumeTrait",
     "MapsTrait",
     "MapContentTrait",
@@ -46,6 +82,9 @@
     "FlowLedStatusTrait",
     "LedStatusTrait",
     "ValleyElectricityTimerTrait",
+    "DustCollectionModeTrait",
+    "WashTowelModeTrait",
+    "SmartWashParamsTrait",
 ]
 
 
@@ -61,6 +100,7 @@ class PropertiesApi(Trait):
     command: CommandTrait
     dnd: DoNotDisturbTrait
     clean_summary: CleanSummaryTrait
+    clean_record: CleanRecordTrait
     sound_volume: SoundVolumeTrait
     rooms: RoomsTrait
     maps: MapsTrait
@@ -74,6 +114,9 @@ class PropertiesApi(Trait):
     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
 
     def __init__(
         self,
@@ -89,8 +132,12 @@ def __init__(
         self._rpc_channel = rpc_channel
         self._mqtt_rpc_channel = mqtt_rpc_channel
         self._map_rpc_channel = map_rpc_channel
+        self._cache = cache
 
         self.status = StatusTrait(product)
+        self.clean_summary = CleanSummaryTrait()
+        self.clean_record = CleanRecordTrait(self.clean_summary)
+        self.consumables = ConsumableTrait()
         self.rooms = RoomsTrait(home_data)
         self.maps = MapsTrait(self.status)
         self.map_content = MapContentTrait(map_parser_config)
@@ -103,7 +150,7 @@ def __init__(
                 # We exclude optional features and them via discover_features
                 if (union_args := get_args(item.type)) is None or len(union_args) > 0:
                     continue
-                _LOGGER.debug("Initializing trait %s", item.name)
+                _LOGGER.debug("Trait '%s' is supported, initializing", item.name)
                 trait = item.type()
                 setattr(self, item.name, trait)
             # This is a hack to allow setting the rpc_channel on all traits. This is
@@ -124,8 +171,12 @@ def _get_rpc_channel(self, trait: V1TraitMixin) -> V1RpcChannel:
 
     async def discover_features(self) -> None:
         """Populate any supported traits that were not initialized in __init__."""
+        _LOGGER.debug("Starting optional trait discovery")
         await self.device_features.refresh()
+        # Dock type also acts like a device feature for some traits.
+        dock_type = await self._dock_type()
 
+        # Dynamically create any traits that need to be populated
         for item in fields(self):
             if (trait := getattr(self, item.name, None)) is not None:
                 continue
@@ -136,22 +187,65 @@ 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]
-            trait = item_type()
-            if not hasattr(trait, "requires_feature"):
-                _LOGGER.debug("Trait missing required feature %s", item.name)
-                continue
-            _LOGGER.debug("Checking for feature %s", trait.requires_feature)
-            is_supported = getattr(self.device_features, trait.requires_feature)
-            # _LOGGER.debug("Device features: %s", self.device_features)
-            if is_supported is None:
-                raise ValueError(f"Device feature '{trait.requires_feature}' on trait '{item.name}' is unknown")
-            if not is_supported:
-                _LOGGER.debug("Disabling optional feature trait %s", item.name)
+            if not self._is_supported(item_type, item.name, dock_type):
+                _LOGGER.debug("Trait '%s' not supported, skipping", item.name)
                 continue
-            _LOGGER.debug("Enabling optional feature trait %s", item.name)
+            _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)
 
+    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_type, "requires_dock_type", None)) is not None:
+            return requires_dock_type(dock_type)
+
+        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:
+            raise ValueError(f"Device feature '{feature_name}' on trait '{name}' is unknown")
+        return is_supported
+
+    async def _dock_type(self) -> RoborockDockTypeCode:
+        """Get the dock type from the status trait or cache."""
+        dock_type = await self._get_cached_trait_data("dock_type")
+        if dock_type is not None:
+            _LOGGER.debug("Using cached dock type: %s", dock_type)
+            try:
+                return RoborockDockTypeCode(dock_type)
+            except ValueError:
+                _LOGGER.debug("Cached dock type %s is invalid, refreshing", dock_type)
+
+        _LOGGER.debug("Starting dock type discovery")
+        await self.status.refresh()
+        _LOGGER.debug("Fetched dock type: %s", self.status.dock_type)
+        if self.status.dock_type is None:
+            # Explicitly set so we reuse cached value next type
+            dock_type = RoborockDockTypeCode.no_dock
+        else:
+            dock_type = self.status.dock_type
+        await self._set_cached_trait_data("dock_type", dock_type)
+        return dock_type
+
+    async def _get_cached_trait_data(self, name: str) -> Any:
+        """Get the dock type from the status trait or cache."""
+        cache_data = await self._cache.get()
+        if cache_data.trait_data is None:
+            cache_data.trait_data = {}
+        _LOGGER.debug("Cached trait data: %s", cache_data.trait_data)
+        return cache_data.trait_data.get(name)
+
+    async def _set_cached_trait_data(self, name: str, value: Any) -> None:
+        """Set trait-specific cached data."""
+        cache_data = await self._cache.get()
+        if cache_data.trait_data is None:
+            cache_data.trait_data = {}
+        cache_data.trait_data[name] = value
+        _LOGGER.debug("Updating cached trait data: %s", cache_data.trait_data)
+        await self._cache.set(cache_data)
+
 
 def create(
     product: HomeDataProduct,