diff --git a/roborock/cli.py b/roborock/cli.py index acc17e69..d98582df 100644 --- a/roborock/cli.py +++ b/roborock/cli.py @@ -21,6 +21,7 @@ roborock> status --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(): diff --git a/roborock/devices/traits/v1/__init__.py b/roborock/devices/traits/v1/__init__.py index d152b01a..7edf5f43 100644 --- a/roborock/devices/traits/v1/__init__.py +++ b/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 diff --git a/roborock/devices/traits/v1/consumeable.py b/roborock/devices/traits/v1/consumeable.py new file mode 100644 index 00000000..c5ffdc1d --- /dev/null +++ b/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]) diff --git a/tests/devices/traits/v1/test_consumable.py b/tests/devices/traits/v1/test_consumable.py new file mode 100644 index 00000000..388d3504 --- /dev/null +++ b/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]) + + +# diff --git a/tests/protocols/__snapshots__/test_v1_protocol.ambr b/tests/protocols/__snapshots__/test_v1_protocol.ambr index cc31e868..0b736df6 100644 --- a/tests/protocols/__snapshots__/test_v1_protocol.ambr +++ b/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 # --- diff --git a/tests/protocols/testdata/v1_protocol/get_consumeables.json b/tests/protocols/testdata/v1_protocol/get_consumeables.json new file mode 100644 index 00000000..e79d1a18 --- /dev/null +++ b/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}]}"}}