Add fan speed presets to VacuumInterface (#1405)

* Added Xiaomi Vaccum Mop 2 Ultra and Pro+

* Updated automatic formatting using black

* Removed duplicated supported models

* Extended dreame test to test all supported models

* Added test for invalid dreame model

* Formatting by black

* import isort

* Updated readme with newly supported models

* Added support for Vacuum interface: VaccumDevice and VacuumMiotDevice

* Fixed isort

* Removing unnecessary pass

* Feedback from PR: VaccumDevice renamed to VacuumInterface

* Step2: VacuumInterface no moreinherits Device

* Unit test fixed

* VaccumInterface published as symbol available for the "interface" package

* Added two methods into VacuumInterface:
- fan_speed_presets()
- set_fan_speed_preset(speed)

* Added vacuum unit test

* Fixed python 3.10

* Imporved test coverage

* Additional increase of test coverage

* Added test for param validation in set_fan_speed_preset()

* unit test simplification

* Feedback from PR

* Minor DOCS improvement

* Minor docs update

* Feedback from PR

* Updated noqa codes for new flake8 version

* remove copyright notice update

* Rework docstrings

Co-authored-by: Teemu Rytilahti <tpr@iki.fi>

devtools/containers.py

@@ -42,7 +42,7 @@ class ModelMapping(DataClassJsonMixin):
     def urn_for_model(self, model: str):
         matches = [inst for inst in self.instances if inst.model == model]
         if len(matches) > 1:
-            print(  # noqa: T001
+            print(  # noqa: T201
                 "WARNING more than a single match for model %s, using the first one: %s"
                 % (model, matches)
             )

devtools/parse_pcap.py

@@ -68,7 +68,7 @@ def read_payloads_from_file(file, tokens: list[str]):
 
         yield src_addr, dst_addr, payload
 
-    print(stats)  # noqa: T001
+    print(stats)  # noqa: T201
 
 
 @app.command()
@@ -77,7 +77,7 @@ def read_file(
 ):
     """Read PCAP file and output decrypted miio communication."""
     for src_addr, dst_addr, payload in read_payloads_from_file(file, token):
-        print(f"{src_addr:<15} -> {dst_addr:<15} {payload}")  # noqa: T001
+        print(f"{src_addr:<15} -> {dst_addr:<15} {payload}")  # noqa: T201
 
 
 if __name__ == "__main__":

miio/integrations/vacuum/dreame/dreamevacuum_miot.py

@@ -8,7 +8,7 @@
 
 from miio.click_common import command, format_output
 from miio.exceptions import DeviceException
-from miio.interfaces import VacuumInterface
+from miio.interfaces import FanspeedPresets, VacuumInterface
 from miio.miot_device import DeviceStatus as DeviceStatusContainer
 from miio.miot_device import MiotDevice, MiotMapping
 
@@ -527,13 +527,22 @@ def set_fan_speed(self, speed: int):
         return self.set_property("cleaning_mode", fanspeed.value)
 
     @command()
-    def fan_speed_presets(self) -> Dict[str, int]:
-        """Return dictionary containing supported fan speeds."""
+    def fan_speed_presets(self) -> FanspeedPresets:
+        """Return available fan speed presets."""
         fanspeeds_enum = _get_cleaning_mode_enum_class(self.model)
         if not fanspeeds_enum:
             return {}
         return _enum_as_dict(fanspeeds_enum)
 
+    @command(click.argument("speed", type=int))
+    def set_fan_speed_preset(self, speed_preset: int) -> None:
+        """Set fan speed preset speed."""
+        if speed_preset not in self.fan_speed_presets().values():
+            raise ValueError(
+                f"Invalid preset speed {speed_preset}, not in: {self.fan_speed_presets().values()}"
+            )
+        self.set_fan_speed(speed_preset)
+
     @command()
     def waterflow(self):
         """Get water flow setting."""

miio/integrations/vacuum/dreame/tests/test_dreamevacuum_miot.py

@@ -176,6 +176,10 @@ def test_fan_speed(self):
         value = self.device.fan_speed()
         assert value == {"Medium": 2}
 
+    def test_set_fan_speed_preset(self):
+        for speed in self.device.fan_speed_presets().values():
+            self.device.set_fan_speed_preset(speed)
+
 
 @pytest.mark.usefixtures("dummydreamef9vacuum")
 class TestDreameF9Vacuum(TestCase):
@@ -253,10 +257,7 @@ def test_waterflow(self):
 
 @pytest.mark.parametrize("model", MIOT_MAPPING.keys())
 def test_dreame_models(model: str):
-    vac = DreameVacuum(model=model)
-    # test _get_cleaning_mode_enum_class returns non-empty mapping
-    fp = vac.fan_speed_presets()
-    assert (fp is not None) and (len(fp) > 0)
+    DreameVacuum(model=model)
 
 
 def test_invalid_dreame_model():

miio/integrations/vacuum/mijia/g1vacuum.py

@@ -5,7 +5,7 @@
 import click
 
 from miio.click_common import EnumType, command, format_output
-from miio.interfaces import VacuumInterface
+from miio.interfaces import FanspeedPresets, VacuumInterface
 from miio.miot_device import DeviceStatus, MiotDevice
 
 _LOGGER = logging.getLogger(__name__)
@@ -373,3 +373,17 @@ def consumable_reset(self, consumable: G1Consumable):
     def set_fan_speed(self, fan_speed: G1FanSpeed):
         """Set fan speed."""
         return self.set_property("fan_speed", fan_speed.value)
+
+    @command()
+    def fan_speed_presets(self) -> FanspeedPresets:
+        """Return available fan speed presets."""
+        return {x.name: x.value for x in G1FanSpeed}
+
+    @command(click.argument("speed", type=int))
+    def set_fan_speed_preset(self, speed_preset: int) -> None:
+        """Set fan speed preset speed."""
+        if speed_preset not in self.fan_speed_presets().values():
+            raise ValueError(
+                f"Invalid preset speed {speed_preset}, not in: {self.fan_speed_presets().values()}"
+            )
+        return self.set_property("fan_speed", speed_preset)

miio/integrations/vacuum/roborock/vacuum.py

@@ -7,7 +7,7 @@
 import os
 import pathlib
 import time
-from typing import Dict, List, Optional, Type, Union
+from typing import List, Optional, Type, Union
 
 import click
 import pytz
@@ -22,7 +22,7 @@
 )
 from miio.device import Device, DeviceInfo
 from miio.exceptions import DeviceException, DeviceInfoUnavailableException
-from miio.interfaces import VacuumInterface
+from miio.interfaces import FanspeedPresets, VacuumInterface
 
 from .vacuumcontainers import (
     CarpetModeStatus,
@@ -619,8 +619,8 @@ def fan_speed(self):
         return self.send("get_custom_mode")[0]
 
     @command()
-    def fan_speed_presets(self) -> Dict[str, int]:
-        """Return dictionary containing supported fan speeds."""
+    def fan_speed_presets(self) -> FanspeedPresets:
+        """Return available fan speed presets."""
 
         def _enum_as_dict(cls):
             return {x.name: x.value for x in list(cls)}
@@ -652,6 +652,15 @@ def _enum_as_dict(cls):
 
         return _enum_as_dict(fanspeeds)
 
+    @command(click.argument("speed", type=int))
+    def set_fan_speed_preset(self, speed_preset: int) -> None:
+        """Set fan speed preset speed."""
+        if speed_preset not in self.fan_speed_presets().values():
+            raise ValueError(
+                f"Invalid preset speed {speed_preset}, not in: {self.fan_speed_presets().values()}"
+            )
+        return self.send("set_custom_mode", [speed_preset])
+
     @command()
     def sound_info(self):
         """Get voice settings."""

miio/integrations/vacuum/roidmi/roidmivacuum_miot.py

@@ -11,7 +11,7 @@
 
 from miio.click_common import EnumType, command
 from miio.integrations.vacuum.roborock.vacuumcontainers import DNDStatus
-from miio.interfaces import VacuumInterface
+from miio.interfaces import FanspeedPresets, VacuumInterface
 from miio.miot_device import DeviceStatus, MiotDevice, MiotMapping
 
 _LOGGER = logging.getLogger(__name__)
@@ -637,6 +637,20 @@ def set_fanspeed(self, fanspeed_mode: FanSpeed):
         """Set fan speed."""
         return self.set_property("fanspeed_mode", fanspeed_mode.value)
 
+    @command()
+    def fan_speed_presets(self) -> FanspeedPresets:
+        """Return available fan speed presets."""
+        return {"Sweep": 0, "Silent": 1, "Basic": 2, "Strong": 3, "FullSpeed": 4}
+
+    @command(click.argument("speed", type=int))
+    def set_fan_speed_preset(self, speed_preset: int) -> None:
+        """Set fan speed preset speed."""
+        if speed_preset not in self.fan_speed_presets().values():
+            raise ValueError(
+                f"Invalid preset speed {speed_preset}, not in: {self.fan_speed_presets().values()}"
+            )
+        return self.set_property("fanspeed_mode", speed_preset)
+
     @command(click.argument("sweep_type", type=EnumType(SweepType)))
     def set_sweep_type(self, sweep_type: SweepType):
         """Set sweep_type."""

miio/integrations/vacuum/roidmi/tests/test_roidmivacuum_miot.py

@@ -182,6 +182,10 @@ def test_parse_forbid_mode2(self):
         )
         assert str(status._parse_forbid_mode(value)) == str(expected_value)
 
+    def test_set_fan_speed_preset(self):
+        for speed in self.device.fan_speed_presets().values():
+            self.device.set_fan_speed_preset(speed)
+
 
 class DummyRoidmiVacuumMiot2(DummyMiotDevice, RoidmiVacuumMiot):
     def __init__(self, *args, **kwargs):

miio/integrations/vacuum/viomi/viomivacuum.py

@@ -58,7 +58,7 @@
     ConsumableStatus,
     DNDStatus,
 )
-from miio.interfaces import VacuumInterface
+from miio.interfaces import FanspeedPresets, VacuumInterface
 from miio.utils import pretty_seconds
 
 _LOGGER = logging.getLogger(__name__)
@@ -338,11 +338,6 @@ def fanspeed(self) -> ViomiVacuumSpeed:
         """Current fan speed."""
         return ViomiVacuumSpeed(self.data["suction_grade"])
 
-    @command()
-    def fan_speed_presets(self) -> Dict[str, int]:
-        """Return dictionary containing supported fanspeeds."""
-        return {x.name: x.value for x in list(ViomiVacuumSpeed)}
-
     @property
     def water_grade(self) -> ViomiWaterGrade:
         """Water grade."""
@@ -677,6 +672,20 @@ def set_fan_speed(self, speed: ViomiVacuumSpeed):
         """Set fanspeed [silent, standard, medium, turbo]."""
         self.send("set_suction", [speed.value])
 
+    @command()
+    def fan_speed_presets(self) -> FanspeedPresets:
+        """Return available fan speed presets."""
+        return {x.name: x.value for x in list(ViomiVacuumSpeed)}
+
+    @command(click.argument("speed", type=int))
+    def set_fan_speed_preset(self, speed_preset: int) -> None:
+        """Set fan speed preset speed."""
+        if speed_preset not in self.fan_speed_presets().values():
+            raise ValueError(
+                f"Invalid preset speed {speed_preset}, not in: {self.fan_speed_presets().values()}"
+            )
+        self.send("set_suction", [speed_preset])
+
     @command(click.argument("watergrade", type=EnumType(ViomiWaterGrade)))
     def set_water_grade(self, watergrade: ViomiWaterGrade):
         """Set water grade.

miio/interfaces/__init__.py

@@ -1,5 +1,5 @@
 """Interfaces API."""
 
-from .vacuuminterface import VacuumInterface
+from .vacuuminterface import FanspeedPresets, VacuumInterface
 
-__all__ = ["VacuumInterface"]
+__all__ = ["FanspeedPresets", "VacuumInterface"]

miio/interfaces/vacuuminterface.py

@@ -1,23 +1,46 @@
 """`VacuumInterface` is an interface (abstract class) with shared API for all vacuum
 devices."""
 from abc import abstractmethod
+from typing import Dict
+
+# Dictionary of predefined fan speeds
+FanspeedPresets = Dict[str, int]
 
 
 class VacuumInterface:
     """Vacuum API interface."""
 
     @abstractmethod
     def home(self):
-        """Return to home."""
+        """Return vacuum robot to home station/dock."""
 
     @abstractmethod
     def start(self):
         """Start cleaning."""
 
     @abstractmethod
     def stop(self):
-        """Validate that Stop cleaning."""
+        """Stop cleaning."""
 
     def pause(self):
-        """Pause cleaning."""
+        """Pause cleaning.
+
+        :raises RuntimeError: if the method is not supported by the device
+        """
         raise RuntimeError("`pause` not supported")
+
+    @abstractmethod
+    def fan_speed_presets(self) -> FanspeedPresets:
+        """Return available fan speed presets.
+
+        The returned object is a dictionary where the key is user-readable name and the
+        value is input for :func:`set_fan_speed_preset()`.
+        """
+
+    @abstractmethod
+    def set_fan_speed_preset(self, speed_preset: int) -> None:
+        """Set fan speed preset speed.
+
+        :param speed_preset: a value from :func:`fan_speed_presets()`
+        :raises ValueError: for invalid preset value
+        """

miio/tests/test_vacuums.py

@@ -0,0 +1,55 @@
+"""Test of vacuum devices."""
+from collections.abc import Iterable
+from typing import List, Sequence, Tuple, Type
+
+import pytest
+
+from miio.device import Device
+from miio.integrations.vacuum.roborock.vacuum import ROCKROBO_V1
+from miio.interfaces import VacuumInterface
+
+# list of all supported vacuum classes
+VACUUM_CLASSES: Tuple[Type[VacuumInterface], ...] = tuple(
+    cl for cl in VacuumInterface.__subclasses__()  # type: ignore
+)
+
+
+def _all_vacuum_models() -> Sequence[Tuple[Type[Device], str]]:
+    """:return: list of tuples with supported vacuum models with corresponding class"""
+    result: List[Tuple[Type[Device], str]] = []
+    for cls in VACUUM_CLASSES:
+        assert issubclass(cls, Device)
+        vacuum_models = cls.supported_models
+        assert isinstance(vacuum_models, Iterable)
+        for model in vacuum_models:
+            result.append((cls, model))
+    return result  # type: ignore
+
+
+@pytest.mark.parametrize("cls, model", _all_vacuum_models())
+def test_vacuum_fan_speed_presets(cls: Type[Device], model: str) -> None:
+    """Test method VacuumInterface.fan_speed_presets()"""
+    if model == ROCKROBO_V1:
+        return  # this model cannot be tested because presets depends on firmware
+    dev = cls("127.0.0.1", "68ffffffffffffffffffffffffffffff", model=model)
+    assert isinstance(dev, VacuumInterface)
+    presets = dev.fan_speed_presets()
+    assert presets is not None, "presets must be defined"
+    assert bool(presets), "presets cannot be empty"
+    assert isinstance(presets, dict), "presets must be dictionary"
+    for name, value in presets.items():
+        assert isinstance(name, str), "presets key must be string"
+        assert name, "presets key cannot be empty"
+        assert isinstance(value, int), "presets value must be integer"
+        assert value >= 0, "presets value must be >= 0"
+
+
+@pytest.mark.parametrize("cls, model", _all_vacuum_models())
+def test_vacuum_set_fan_speed_presets_fails(cls: Type[Device], model: str) -> None:
+    """Test method VacuumInterface.fan_speed_presets()"""
+    if model == ROCKROBO_V1:
+        return  # this model cannot be tested because presets depends on firmware
+    dev = cls("127.0.0.1", "68ffffffffffffffffffffffffffffff", model=model)
+    assert isinstance(dev, VacuumInterface)
+    with pytest.raises(ValueError):
+        dev.set_fan_speed_preset(-1)