Add consumables module for vacuums (#1327)

Co-authored-by: Steven B <51370195+sdb9696@users.noreply.github.com>

kasa/cli/vacuum.py

@@ -51,3 +51,34 @@ async def records_list(dev: Device) -> None:
             f"* {record.timestamp}: cleaned {record.clean_area} {rec.area_unit}"
             f" in {record.clean_time}"
         )
+
+
+@vacuum.group(invoke_without_command=True, name="consumables")
+@pass_dev_or_child
+@click.pass_context
+async def consumables(ctx: click.Context, dev: Device) -> None:
+    """List device consumables."""
+    if not (cons := dev.modules.get(Module.Consumables)):
+        error("This device does not support consumables.")
+
+    if not ctx.invoked_subcommand:
+        for c in cons.consumables.values():
+            click.echo(f"{c.name} ({c.id}): {c.used} used, {c.remaining} remaining")
+
+
+@consumables.command(name="reset")
+@click.argument("consumable_id", required=True)
+@pass_dev_or_child
+async def reset_consumable(dev: Device, consumable_id: str) -> None:
+    """Reset the consumable used/remaining time."""
+    cons = dev.modules[Module.Consumables]
+
+    if consumable_id not in cons.consumables:
+        error(
+            f"Consumable {consumable_id} not found in "
+            f"device consumables: {', '.join(cons.consumables.keys())}."
+        )
+
+    await cons.reset_consumable(consumable_id)
+
+    click.echo(f"Consumable {consumable_id} reset")

kasa/module.py

@@ -165,6 +165,7 @@ class Module(ABC):
 
     # Vacuum modules
     Clean: Final[ModuleName[smart.Clean]] = ModuleName("Clean")
+    Consumables: Final[ModuleName[smart.Consumables]] = ModuleName("Consumables")
     Dustbin: Final[ModuleName[smart.Dustbin]] = ModuleName("Dustbin")
     Speaker: Final[ModuleName[smart.Speaker]] = ModuleName("Speaker")
     Mop: Final[ModuleName[smart.Mop]] = ModuleName("Mop")

kasa/smart/modules/__init__.py

@@ -14,6 +14,7 @@ from .cleanrecords import CleanRecords
 from .cloud import Cloud
 from .color import Color
 from .colortemperature import ColorTemperature
+from .consumables import Consumables
 from .contactsensor import ContactSensor
 from .devicemodule import DeviceModule
 from .dustbin import Dustbin
@@ -76,6 +77,7 @@ __all__ = [
     "FrostProtection",
     "Thermostat",
     "Clean",
+    "Consumables",
     "CleanRecords",
     "SmartLightEffect",
     "OverheatProtection",

kasa/smart/modules/consumables.py

@@ -0,0 +1,170 @@
+"""Implementation of vacuum consumables."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Mapping
+from dataclasses import dataclass
+from datetime import timedelta
+
+from ...feature import Feature
+from ..smartmodule import SmartModule
+
+_LOGGER = logging.getLogger(__name__)
+
+
+@dataclass
+class _ConsumableMeta:
+    """Consumable meta container."""
+
+    #: Name of the consumable.
+    name: str
+    #: Internal id of the consumable
+    id: str
+    #: Data key in the device reported data
+    data_key: str
+    #: Lifetime
+    lifetime: timedelta
+
+
+@dataclass
+class Consumable:
+    """Consumable container."""
+
+    #: Name of the consumable.
+    name: str
+    #: Id of the consumable
+    id: str
+    #: Lifetime
+    lifetime: timedelta
+    #: Used
+    used: timedelta
+    #: Remaining
+    remaining: timedelta
+    #: Device data key
+    _data_key: str
+
+
+CONSUMABLE_METAS = [
+    _ConsumableMeta(
+        "Main brush",
+        id="main_brush",
+        data_key="roll_brush_time",
+        lifetime=timedelta(hours=400),
+    ),
+    _ConsumableMeta(
+        "Side brush",
+        id="side_brush",
+        data_key="edge_brush_time",
+        lifetime=timedelta(hours=200),
+    ),
+    _ConsumableMeta(
+        "Filter",
+        id="filter",
+        data_key="filter_time",
+        lifetime=timedelta(hours=200),
+    ),
+    _ConsumableMeta(
+        "Sensor",
+        id="sensor",
+        data_key="sensor_time",
+        lifetime=timedelta(hours=30),
+    ),
+    _ConsumableMeta(
+        "Charging contacts",
+        id="charging_contacts",
+        data_key="charge_contact_time",
+        lifetime=timedelta(hours=30),
+    ),
+    # Unknown keys: main_brush_lid_time, rag_time
+]
+
+
+class Consumables(SmartModule):
+    """Implementation of vacuum consumables."""
+
+    REQUIRED_COMPONENT = "consumables"
+    QUERY_GETTER_NAME = "getConsumablesInfo"
+
+    _consumables: dict[str, Consumable] = {}
+
+    def _initialize_features(self) -> None:
+        """Initialize features."""
+        for c_meta in CONSUMABLE_METAS:
+            if c_meta.data_key not in self.data:
+                continue
+
+            self._add_feature(
+                Feature(
+                    self._device,
+                    id=f"{c_meta.id}_used",
+                    name=f"{c_meta.name} used",
+                    container=self,
+                    attribute_getter=lambda _, c_id=c_meta.id: self._consumables[
+                        c_id
+                    ].used,
+                    category=Feature.Category.Debug,
+                    type=Feature.Type.Sensor,
+                )
+            )
+
+            self._add_feature(
+                Feature(
+                    self._device,
+                    id=f"{c_meta.id}_remaining",
+                    name=f"{c_meta.name} remaining",
+                    container=self,
+                    attribute_getter=lambda _, c_id=c_meta.id: self._consumables[
+                        c_id
+                    ].remaining,
+                    category=Feature.Category.Info,
+                    type=Feature.Type.Sensor,
+                )
+            )
+
+            self._add_feature(
+                Feature(
+                    self._device,
+                    id=f"{c_meta.id}_reset",
+                    name=f"Reset {c_meta.name.lower()} consumable",
+                    container=self,
+                    attribute_setter=lambda c_id=c_meta.id: self.reset_consumable(c_id),
+                    category=Feature.Category.Debug,
+                    type=Feature.Type.Action,
+                )
+            )
+
+    async def _post_update_hook(self) -> None:
+        """Update the consumables."""
+        if not self._consumables:
+            for consumable_meta in CONSUMABLE_METAS:
+                if consumable_meta.data_key not in self.data:
+                    continue
+                used = timedelta(minutes=self.data[consumable_meta.data_key])
+                consumable = Consumable(
+                    id=consumable_meta.id,
+                    name=consumable_meta.name,
+                    lifetime=consumable_meta.lifetime,
+                    used=used,
+                    remaining=consumable_meta.lifetime - used,
+                    _data_key=consumable_meta.data_key,
+                )
+                self._consumables[consumable_meta.id] = consumable
+        else:
+            for consumable in self._consumables.values():
+                consumable.used = timedelta(minutes=self.data[consumable._data_key])
+                consumable.remaining = consumable.lifetime - consumable.used
+
+    async def reset_consumable(self, consumable_id: str) -> dict:
+        """Reset consumable stats."""
+        consumable_name = self._consumables[consumable_id]._data_key.removesuffix(
+            "_time"
+        )
+        return await self.call(
+            "resetConsumablesTime", {"reset_list": [consumable_name]}
+        )
+
+    @property
+    def consumables(self) -> Mapping[str, Consumable]:
+        """Get list of consumables on the device."""
+        return self._consumables

tests/cli/test_vacuum.py

@@ -45,6 +45,49 @@ async def test_vacuum_records_list(dev, mocker: MockerFixture, runner):
     assert res.exit_code == 0
 
 
+@vacuum_devices
+async def test_vacuum_consumables(dev, runner):
+    """Test that vacuum consumables calls the expected methods."""
+    cons = dev.modules.get(Module.Consumables)
+    assert cons
+
+    res = await runner.invoke(vacuum, ["consumables"], obj=dev, catch_exceptions=False)
+
+    expected = ""
+    for c in cons.consumables.values():
+        expected += f"{c.name} ({c.id}): {c.used} used, {c.remaining} remaining\n"
+
+    assert expected in res.output
+    assert res.exit_code == 0
+
+
+@vacuum_devices
+async def test_vacuum_consumables_reset(dev, mocker: MockerFixture, runner):
+    """Test that vacuum consumables reset calls the expected methods."""
+    cons = dev.modules.get(Module.Consumables)
+    assert cons
+
+    reset_consumable_mock = mocker.spy(cons, "reset_consumable")
+    for c_id in cons.consumables:
+        reset_consumable_mock.reset_mock()
+        res = await runner.invoke(
+            vacuum, ["consumables", "reset", c_id], obj=dev, catch_exceptions=False
+        )
+        reset_consumable_mock.assert_awaited_once_with(c_id)
+        assert f"Consumable {c_id} reset" in res.output
+        assert res.exit_code == 0
+
+    res = await runner.invoke(
+        vacuum, ["consumables", "reset", "foobar"], obj=dev, catch_exceptions=False
+    )
+    expected = (
+        "Consumable foobar not found in "
+        f"device consumables: {', '.join(cons.consumables.keys())}."
+    )
+    assert expected in res.output.replace("\n", "")
+    assert res.exit_code != 0
+
+
 @plug_iot
 async def test_non_vacuum(dev, mocker: MockerFixture, runner):
     """Test that vacuum commands return an error if executed on a non-vacuum."""
@@ -59,3 +102,13 @@ async def test_non_vacuum(dev, mocker: MockerFixture, runner):
     )
     assert "This device does not support records" in res.output
     assert res.exit_code != 0
+
+    res = await runner.invoke(vacuum, ["consumables"], obj=dev, catch_exceptions=False)
+    assert "This device does not support consumables" in res.output
+    assert res.exit_code != 0
+
+    res = await runner.invoke(
+        vacuum, ["consumables", "reset", "foobar"], obj=dev, catch_exceptions=False
+    )
+    assert "This device does not support consumables" in res.output
+    assert res.exit_code != 0

tests/fakeprotocol_smart.py

@@ -687,6 +687,7 @@ class FakeSmartTransport(BaseTransport):
             "add_child_device_list",  # hub pairing
             "remove_child_device_list",  # hub pairing
             "playSelectAudio",  # vacuum special actions
+            "resetConsumablesTime",  # vacuum special actions
         ]:
             return {"error_code": 0}
         elif method[:3] == "set":

tests/smart/modules/test_consumables.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from datetime import timedelta
+
+import pytest
+from pytest_mock import MockerFixture
+
+from kasa import Module
+from kasa.smart import SmartDevice
+from kasa.smart.modules.consumables import CONSUMABLE_METAS
+
+from ...device_fixtures import get_parent_and_child_modules, parametrize
+
+consumables = parametrize(
+    "has consumables", component_filter="consumables", protocol_filter={"SMART"}
+)
+
+
+@consumables
+@pytest.mark.parametrize(
+    "consumable_name", [consumable.id for consumable in CONSUMABLE_METAS]
+)
+@pytest.mark.parametrize("postfix", ["used", "remaining"])
+async def test_features(dev: SmartDevice, consumable_name: str, postfix: str):
+    """Test that features are registered and work as expected."""
+    consumables = next(get_parent_and_child_modules(dev, Module.Consumables))
+    assert consumables is not None
+
+    feature_name = f"{consumable_name}_{postfix}"
+
+    feat = consumables._device.features[feature_name]
+    assert isinstance(feat.value, timedelta)
+
+
+@consumables
+@pytest.mark.parametrize(
+    ("consumable_name", "data_key"),
+    [(consumable.id, consumable.data_key) for consumable in CONSUMABLE_METAS],
+)
+async def test_erase(
+    dev: SmartDevice, mocker: MockerFixture, consumable_name: str, data_key: str
+):
+    """Test autocollection switch."""
+    consumables = next(get_parent_and_child_modules(dev, Module.Consumables))
+    call = mocker.spy(consumables, "call")
+
+    feature_name = f"{consumable_name}_reset"
+    feat = dev._features[feature_name]
+    await feat.set_value(True)
+
+    call.assert_called_with(
+        "resetConsumablesTime", {"reset_list": [data_key.removesuffix("_time")]}
+    )