feat: add support for getting and reseting consumables (#502)

* feat: add support for getting and reseting consumables

* chore: add a class comment about availability

* chore: remove whitespace

Author: allenporter

roborock/cli.py

@@ -21,6 +21,7 @@
 roborock> status --device_id <device_id>
 ```
 """
+
 import asyncio
 import datetime
 import functools
@@ -46,6 +47,7 @@
 from roborock.devices.device_manager import DeviceManager, create_device_manager, create_home_data_api
 from roborock.devices.traits import Trait
 from roborock.devices.traits.v1 import V1TraitMixin
+from roborock.devices.traits.v1.consumeable import ConsumableAttribute
 from roborock.protocol import MessageParser
 from roborock.version_1_apis.roborock_mqtt_client_v1 import RoborockMqttClientV1
 from roborock.web_api import RoborockApiClient
@@ -449,6 +451,30 @@ async def maps(ctx, device_id: str):
     await _display_v1_trait(context, device_id, lambda v1: v1.maps)
 
 
+@session.command()
+@click.option("--device_id", required=True)
+@click.pass_context
+@async_command
+async def consumables(ctx, device_id: str):
+    """Get device consumables."""
+    context: RoborockContext = ctx.obj
+    await _display_v1_trait(context, device_id, lambda v1: v1.consumables)
+
+
+@session.command()
+@click.option("--device_id", required=True)
+@click.option("--consumable", required=True, type=click.Choice([e.value for e in ConsumableAttribute]))
+@click.pass_context
+@async_command
+async def reset_consumable(ctx, device_id: str, consumable: str):
+    """Reset a specific consumable attribute."""
+    context: RoborockContext = ctx.obj
+    trait = await _v1_trait(context, device_id, lambda v1: v1.consumables)
+    attribute = ConsumableAttribute.from_str(consumable)
+    await trait.reset_consumable(attribute)
+    click.echo(f"Reset {consumable} for device {device_id}")
+
+
 @click.command()
 @click.option("--device_id", required=True)
 @click.option("--cmd", required=True)
@@ -691,6 +717,8 @@ def write_markdown_table(product_features: dict[str, dict[str, any]], all_featur
 cli.add_command(volume)
 cli.add_command(set_volume)
 cli.add_command(maps)
+cli.add_command(consumables)
+cli.add_command(reset_consumable)
 
 
 def main():

roborock/devices/traits/v1/__init__.py

@@ -9,6 +9,7 @@
 
 from .clean_summary import CleanSummaryTrait
 from .common import V1TraitMixin
+from .consumeable import ConsumableTrait
 from .do_not_disturb import DoNotDisturbTrait
 from .maps import MapsTrait
 from .status import StatusTrait
@@ -24,6 +25,7 @@
     "CleanSummaryTrait",
     "SoundVolumeTrait",
     "MapsTrait",
+    "ConsumableTrait",
 ]
 
 
@@ -40,6 +42,7 @@ class PropertiesApi(Trait):
     clean_summary: CleanSummaryTrait
     sound_volume: SoundVolumeTrait
     maps: MapsTrait
+    consumables: ConsumableTrait
 
     # In the future optional fields can be added below based on supported features
 

roborock/devices/traits/v1/consumeable.py

@@ -0,0 +1,47 @@
+"""Trait for managing consumable attributes.
+
+A consumable attribute is one that is expected to be replaced or refilled
+periodically, such as filters, brushes, etc.
+"""
+
+from enum import StrEnum
+from typing import Self
+
+from roborock.containers import Consumable
+from roborock.devices.traits.v1 import common
+from roborock.roborock_typing import RoborockCommand
+
+__all__ = [
+    "ConsumableTrait",
+]
+
+
+class ConsumableAttribute(StrEnum):
+    """Enum for consumable attributes."""
+
+    SENSOR_DIRTY_TIME = "sensor_dirty_time"
+    FILTER_WORK_TIME = "filter_work_time"
+    SIDE_BRUSH_WORK_TIME = "side_brush_work_time"
+    MAIN_BRUSH_WORK_TIME = "main_brush_work_time"
+
+    @classmethod
+    def from_str(cls, value: str) -> Self:
+        """Create a ConsumableAttribute from a string value."""
+        for member in cls:
+            if member.value == value:
+                return member
+        raise ValueError(f"Unknown ConsumableAttribute: {value}")
+
+
+class ConsumableTrait(Consumable, common.V1TraitMixin):
+    """Trait for managing consumable attributes on Roborock devices.
+
+    After the first refresh, you can tell what consumables are supported by
+    checking which attributes are not None.
+    """
+
+    command = RoborockCommand.GET_CONSUMABLE
+
+    async def reset_consumable(self, consumable: ConsumableAttribute) -> None:
+        """Reset a specific consumable attribute on the device."""
+        await self.rpc_channel.send_command(RoborockCommand.RESET_CONSUMABLE, params=[consumable.value])

tests/devices/traits/v1/test_consumable.py

@@ -0,0 +1,70 @@
+"""Tests for the DoNotDisturbTrait class."""
+
+from unittest.mock import AsyncMock
+
+import pytest
+
+from roborock.devices.device import RoborockDevice
+from roborock.devices.traits.v1.consumeable import ConsumableAttribute, ConsumableTrait
+from roborock.roborock_typing import RoborockCommand
+
+CONSUMABLE_DATA = [
+    {
+        "main_brush_work_time": 879348,
+        "side_brush_work_time": 707618,
+        "filter_work_time": 738722,
+        "filter_element_work_time": 0,
+        "sensor_dirty_time": 455517,
+    }
+]
+
+
+@pytest.fixture
+def consumable_trait(device: RoborockDevice) -> ConsumableTrait:
+    """Create a ConsumableTrait instance with mocked dependencies."""
+    assert device.v1_properties
+    return device.v1_properties.consumables
+
+
+async def test_get_consumable_data_success(consumable_trait: ConsumableTrait, mock_rpc_channel: AsyncMock) -> None:
+    """Test successfully getting consumable data."""
+    # Setup mock to return the sample consumable data
+    mock_rpc_channel.send_command.return_value = CONSUMABLE_DATA
+
+    # Call the method
+    await consumable_trait.refresh()
+    # Verify the result
+    assert consumable_trait.main_brush_work_time == 879348
+    assert consumable_trait.side_brush_work_time == 707618
+    assert consumable_trait.filter_work_time == 738722
+    assert consumable_trait.filter_element_work_time == 0
+    assert consumable_trait.sensor_dirty_time == 455517
+
+    # Verify the RPC call was made correctly
+    mock_rpc_channel.send_command.assert_called_once_with(RoborockCommand.GET_CONSUMABLE)
+
+
+@pytest.mark.parametrize(
+    ("consumable", "reset_param"),
+    [
+        (ConsumableAttribute.MAIN_BRUSH_WORK_TIME, "main_brush_work_time"),
+        (ConsumableAttribute.SIDE_BRUSH_WORK_TIME, "side_brush_work_time"),
+        (ConsumableAttribute.FILTER_WORK_TIME, "filter_work_time"),
+        (ConsumableAttribute.SENSOR_DIRTY_TIME, "sensor_dirty_time"),
+    ],
+)
+async def test_reset_consumable_data(
+    consumable_trait: ConsumableTrait,
+    mock_rpc_channel: AsyncMock,
+    consumable: ConsumableAttribute,
+    reset_param: str,
+) -> None:
+    """Test successfully resetting consumable data."""
+    # Call the method
+    await consumable_trait.reset_consumable(consumable)
+
+    # Verify the RPC call was made correctly with expected parameters
+    mock_rpc_channel.send_command.assert_called_once_with(RoborockCommand.RESET_CONSUMABLE, params=[reset_param])
+
+
+#

tests/protocols/__snapshots__/test_v1_protocol.ambr

@@ -33,6 +33,22 @@
   ]
   '''
 # ---
+# name: test_decode_rpc_payload[get_consumeables]
+  20001
+# ---
+# name: test_decode_rpc_payload[get_consumeables].1
+  '''
+  [
+    {
+      "main_brush_work_time": 879348,
+      "side_brush_work_time": 707618,
+      "filter_work_time": 738722,
+      "filter_element_work_time": 0,
+      "sensor_dirty_time": 455517
+    }
+  ]
+  '''
+# ---
 # name: test_decode_rpc_payload[get_dnd]
   20002
 # ---

tests/protocols/testdata/v1_protocol/get_consumeables.json

@@ -0,0 +1 @@
+{"t":1759038395,"dps":{"102":"{\"id\":20001,\"result\":[{\"main_brush_work_time\":879348,\"side_brush_work_time\":707618,\"filter_work_time\":738722,\"filter_element_work_time\":0,\"sensor_dirty_time\":455517}]}"}}