8000 Add consumables module for vacuums (#1327) · python-kasa/python-kasa@a03a4b1 · GitHub
[go: up one dir, main page]

Skip to content

Commit a03a4b1

Browse files
rytilahtisdb9696
andauthored
Add consumables module for vacuums (#1327)
Co-authored-by: Steven B <51370195+sdb9696@users.noreply.github.com>
1 parent 0508546 commit a03a4b1

File tree

7 files changed

+311
-0
lines changed

7 files changed

+311
-0
lines changed

kasa/cli/vacuum.py

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,3 +51,34 @@ async def records_list(dev: Device) -> None:
5151
f"* {record.timestamp}: cleaned {record.clean_area} {rec.area_unit}"
5252
f" in {record.clean_time}"
5353
)
54+
55+
56+
@vacuum.group(invoke_without_command=True, name="consumables")
57+
@pass_dev_or_child
58+
@click.pass_context
59+
async def consumables(ctx: click.Context, dev: Device) -> None:
60+
"""List device consumables."""
61+
if not (cons := dev.modules.get(Module.Consumables)):
62+
error("This device does not support consumables.")
63+
64+
if not ctx.invoked_subcommand:
65+
for c in cons.consumables.values():
66+
click.echo(f"{c.name} ({c.id}): {c.used} used, {c.remaining} remaining")
67+
68+
69+
@consumables.command(name="reset")
70+
@click.argument("consumable_id", required=True)
71+
@pass_dev_or_child
72+
async def reset_consumable(dev: Device, consumable_id: str) -> None:
73+
"""Reset the consumable used/remaining time."""
74+
cons = dev.modules[Module.Consumables]
75+
76+
if consumable_id not in cons.consumables:
77+
57AE error(
78+
f"Consumable {consumable_id} not found in "
79+
f"device consumables: {', '.join(cons.consumables.keys())}."
80+
)
81+
82+
await cons.reset_consumable(consumable_id)
83+
84+
click.echo(f"Consumable {consumable_id} reset")

kasa/module.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -165,6 +165,7 @@ class Module(ABC):
165165

166166
# Vacuum modules
167167
Clean: Final[ModuleName[smart.Clean]] = ModuleName("Clean")
168+
Consumables: Final[ModuleName[smart.Consumables]] = ModuleName("Consumables")
168169
Dustbin: Final[ModuleName[smart.Dustbin]] = ModuleName("Dustbin")
169170
Speaker: Final[ModuleName[smart.Speaker]] = ModuleName("Speaker")
170171
Mop: Final[ModuleName[smart.Mop]] = ModuleName("Mop")

kasa/smart/modules/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
from .cloud import Cloud
1515
from .color import Color
1616
from .colortemperature import ColorTemperature
17+
from .consumables import Consumables
1718
from .contactsensor import ContactSensor
1819
from .devicemodule import DeviceModule
1920
from .dustbin import Dustbin
@@ -76,6 +77,7 @@
7677
"FrostProtection",
7778
"Thermostat",
7879
"Clean",
80+
"Consumables",
7981
"CleanRecords",
8082
"SmartLightEffect",
8183
"OverheatProtection",

kasa/smart/modules/consumables.py

Lines changed: 170 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,170 @@
1+
"""Implementation of vacuum consumables."""
2+
3+
from __future__ import annotations
4+
5+
import logging
6+
from collections.abc import Mapping
7+
from dataclasses import dataclass
8+
from datetime import timedelta
9+
10+
from ...feature import Feature
11+
from ..smartmodule import SmartModule
12+
13+
_LOGGER = logging.getLogger(__name__)
14+
15+
16+
@dataclass
17+
class _ConsumableMeta:
18+
"""Consumable meta container."""
19+
20+
#: Name of the consumable.
21+
name: str
22+
#: Internal id of the consumable
23+
id: str
24+
#: Data key in the device reported data
25+
data_key: str
26+
#: Lifetime
27+
lifetime: timedelta
28+
29+
30+
@dataclass
31+
class Consumable:
32+
"""Consumable container."""
33+
34+
#: Name of the consumable.
35+
name: str
36+
#: Id of the consumable
37+
id: str
38+
#: Lifetime
39+
lifetime: timedelta
40+
#: Used
41+
used: timedelta
42+
#: Remaining
43+
remaining: timedelta
44+
#: Device data key
45+
_data_key: str
46+
47+
48+
CONSUMABLE_METAS = [
49+
_ConsumableMeta(
50+
"Main brush",
51+
id="main_brush",
52+
data_key="roll_brush_time",
53+
lifetime=timedelta(hours=400),
54+
),
55+
_ConsumableMeta(
56+
"Side brush",
57+
id="side_brush",
58+
data_key="edge_brush_time",
59+
lifetime=timedelta(hours=200),
60+
),
61+
_ConsumableMeta(
62+
"Filter",
63+
id="filter",
64+
data_key="filter_time",
65+
lifetime=timedelta(hours=200),
66+
),
67+
_ConsumableMeta(
68+
"Sensor",
69+
id="sensor",
70+
data_key="sensor_time",
71+
lifetime=timedelta(hours=30),
72+
),
73+
_ConsumableMeta(
74+
"Charging contacts",
75+
id="charging_contacts",
76+
data_key="charge_contact_time",
77+
lifetime=timedelta(hours=30),
78+
),
79+
# Unknown keys: main_brush_lid_time, rag_time
80+
]
81+
82+
83+
class Consumables(SmartModule):
84+
"""Implementation of vacuum consumables."""
85+
86+
REQUIRED_COMPONENT = "consumables"
87+
QUERY_GETTER_NAME = "ge 1C6A tConsumablesInfo"
88+
89+
_consumables: dict[str, Consumable] = {}
90+
91+
def _initialize_features(self) -> None:
92+
"""Initialize features."""
93+
for c_meta in CONSUMABLE_METAS:
94+
if c_meta.data_key not in self.data:
95+
continue
96+
97+
self._add_feature(
98+
Feature(
99+
self._device,
100+
id=f"{c_meta.id}_used",
101+
name=f"{c_meta.name} used",
102+
container=self,
103+
attribute_getter=lambda _, c_id=c_meta.id: self._consumables[
104+
c_id
105+
].used,
106+
category=Feature.Category.Debug,
107+
type=Feature.Type.Sensor,
108+
)
109+
)
110+
111+
self._add_feature(
112+
Feature(
113+
self._device,
114+
id=f"{c_meta.id}_remaining",
115+
name=f"{c_meta.name} remaining",
116+
container=self,
117+
attribute_getter=lambda _, c_id=c_meta.id: self._consumables[
118+
c_id
119+
].remaining,
120+
category=Feature.Category.Info,
121+
type=Feature.Type.Sensor,
122+
)
123+
)
124+
125+
self._add_feature(
126+
Feature(
127+
self._device,
128+
id=f"{c_meta.id}_reset",
129+
name=f"Reset {c_meta.name.lower()} consumable",
130+
container=self,
131+
attribute_setter=lambda c_id=c_meta.id: self.reset_consumable(c_id),
132+
category=Feature.Category.Debug,
133+
type=Feature.Type.Action,
134+
)
135+
)
136+
137+
async def _post_update_hook(self) -> None:
138+
"""Update the consumables."""
139+
if not self._consumables:
140+
for consumable_meta in CONSUMABLE_METAS:
141+
if consumable_meta.data_key not in self.data:
142+
continue
143+
used = timedelta(minutes=self.data[consumable_meta.data_key])
144+
consumable = Consumable(
145+
id=consumable_meta.id,
146+
name=consumable_meta.name,
147+
lifetime=consumable_meta.lifetime,
148+
used=used,
149+
remaining=consumable_meta.lifetime - used,
150+
_data_key=consumable_meta.data_key,
151+
)
152+
self._consumables[consumable_meta.id] = consumable
153+
else:
154+
for consumable in self._consumables.values():
155+
consumable.used = timedelta(minutes=self.data[consumable._data_key])
156+
consumable.remaining = consumable.lifetime - consumable.used
157+
158+
async def reset_consumable(self, consumable_id: str) -> dict:
159+
"""Reset consumable stats."""
160+
consumable_name = self._consumables[consumable_id]._data_key.removesuffix(
161+
"_time"
162+
)
163+
return await self.call(
164+
"resetConsumablesTime", {"reset_list": [consumable_name]}
165+
)
166+
167+
@property
168+
def consumables(self) -> Mapping[str, Consumable]:
169+
"""Get list of consumables on the device."""
170+
return self._consumables

tests/cli/test_vacuum.py

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,49 @@ async def test_vacuum_records_list(dev, mocker: MockerFixture, runner):
4545
assert res.exit_code == 0
4646

4747

48+
@vacuum_devices
49+
async def test_vacuum_consumables(dev, runner):
50+
"""Test that vacuum consumables calls the expected methods."""
51+
cons = dev.modules.get(Module.Consumables)
52+
assert cons
53+
54+
res = await runner.invoke(vacuum, ["consumables"], obj=dev, catch_exceptions=False)
55+
56+
expected = ""
57+
for c in cons.consumables.values():
58+
expected += f"{c.name} ({c.id}): {c.used} used, {c.remaining} remaining\n"
59+
60+
assert expected in res.output
61+
assert res.exit_code == 0
62+
63+
64+
@vacuum_devices
65+
async def test_vacuum_consumables_reset(dev, mocker: MockerFixture, runner):
66+
"""Test that vacuum consumables reset calls the expected methods."""
67+
cons = dev.modules.get(Module.Consumables)
68+
assert cons
69+
70+
reset_consumable_mock = mocker.spy(cons, "reset_consumable")
71+
for c_id in cons.consumables:
72+
reset_consumable_mock.reset_mock()
73+
res = await runner.invoke(
74+
vacuum, ["consumables", "reset", c_id], obj=dev, catch_exceptions=False
75+
)
76+
reset_consumable_mock.assert_awaited_once_with(c_id)
77+
assert f"Consumable {c_id} reset" in res.output
78+
assert res.exit_code == 0
79+
80+
res = await runner.invoke(
81+
vacuum, ["consumables", "reset", "foobar"], obj=dev, catch_exceptions=False
82+
)
83+
expected = (
84+
"Consumable foobar not found in "
85+
f"device consumables: {', '.join(cons.consumables.keys())}."
86+
)
87+
assert expected in res.output.replace("\n", "")
88+
assert res.exit_code != 0
89+
90+
4891
@plug_iot
4992
async def test_non_vacuum(dev, mocker: MockerFixture, runner):
5093
"""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):
59102
)
60103
assert "This device does not support records" in res.output
61104
assert res.exit_code != 0
105+
106+
res = await runner.invoke(vacuum, ["consumables"], obj=dev, catch_exceptions=False)
107+
assert "This device does not support consumables" in res.output
108+
assert res.exit_code != 0
109+
110+
res = await runner.invoke(
111+
vacuum, ["consumables", "reset", "foobar"], obj=dev, catch_exceptions=False
112+
)
113+
assert "This device does not support consumables" in res.output
114+
assert res.exit_code != 0

tests/fakeprotocol_smart.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -687,6 +687,7 @@ async def _send_request(self, request_dict: dict):
687687
"add_child_device_list", # hub pairing
688688
"remove_child_device_list", # hub pairing
689689
"playSelectAudio", # vacuum special actions
690+
"resetConsumablesTime", # vacuum special actions
690691
]:
691692
return {"error_code": 0}
692693
elif method[:3] == "set":
Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
from __future__ import annotations
2+
3+
from datetime import timedelta
4+
5+
import pytest
6+
from pytest_mock import MockerFixture
7+
8+
from kasa import Module
9+
from kasa.smart import SmartDevice
10+
from kasa.smart.modules.consumables import CONSUMABLE_METAS
11+
12+
from ...device_fixtures import get_parent_and_child_modules, parametrize
13+
14+
consumables = parametrize(
15+
"has consumables", component_filter="consumables", protocol_filter={"SMART"}
16+
)
17+
18+
19+
@consumables
20+
@pytest.mark.parametrize(
21+
"consumable_name", [consumable.id for consumable in CONSUMABLE_METAS]
22+
)
23+
@pytest.mark.parametrize("postfix", ["used", "remaining"])
24+
async def test_features(dev: SmartDevice, consumable_name: str, postfix: str):
25+
"""Test that features are registered and work as expected."""
26+
consumables = next(get_parent_and_child_modules(dev, Module.Consumables))
27+
assert consumables is not None
28+
29+
feature_name = f"{consumable_name}_{postfix}"
30+
31+
feat = consumables._device.features[feature_name]
32+
assert isinstance(feat.value, timedelta)
33+
34+
35+
@consumables
36+
@pytest.mark.parametrize(
37+
("consumable_name", "data_key"),
38+
[(consumable.id, consumable.data_key) for consumable in CONSUMABLE_METAS],
39+
)
40+
async def test_erase(
41+
dev: SmartDevice, mocker: MockerFixture, consumable_name: str, data_key: str
42+
):
43+
"""Test autocollection switch."""
44+
consumables = next(get_parent_and_child_modules(dev, Module.Consumables))
45+
call = mocker.spy(consumables, "call")
46+
47+
feature_name = f"{consumable_name}_reset"
48+
feat = dev._features[feature_name]
49+
await feat.set_value(True)
50+
51+
call.assert_called_with(
52+
"resetConsumablesTime", {"reset_list": [data_key.removesuffix("_time")]}
53+
)

0 commit comments

Comments
 (0)
0