Add S7 mop scrub intensity (#1236)

* Add mop scrub intensity

* Add vacuum model check

* Add tests for raising exception

* Add more tests

* Update test method name

miio/integrations/vacuum/roborock/tests/test_vacuum.py

@@ -7,7 +7,13 @@
 from miio import RoborockVacuum, Vacuum, VacuumStatus
 from miio.tests.dummies import DummyDevice
 
-from ..vacuum import CarpetCleaningMode, MopMode
+from ..vacuum import (
+    ROCKROBO_S7,
+    CarpetCleaningMode,
+    MopIntensity,
+    MopMode,
+    VacuumException,
+)
 
 
 class DummyVacuum(DummyDevice, RoborockVacuum):
@@ -312,10 +318,45 @@ def test_mop_mode(self):
         with patch.object(self.device, "send", return_value=[32453]):
             assert self.device.mop_mode() is None
 
+    def test_mop_intensity_model_check(self):
+        """Test Roborock S7 check when getting mop intensity."""
+        with pytest.raises(VacuumException):
+            self.device.mop_intensity()
+
+    def test_set_mop_intensity_model_check(self):
+        """Test Roborock S7 check when setting mop intensity."""
+        with pytest.raises(VacuumException):
+            self.device.set_mop_intensity(MopIntensity.Intense)
+
 
 def test_deprecated_vacuum(caplog):
     with pytest.deprecated_call():
         Vacuum("127.1.1.1", "68ffffffffffffffffffffffffffffff")
 
     with pytest.deprecated_call():
         from miio.vacuum import ROCKROBO_S6  # noqa: F401
+
+
+class DummyVacuumS7(DummyVacuum):
+    def __init__(self, *args, **kwargs):
+        self._model = ROCKROBO_S7
+
+
+@pytest.fixture(scope="class")
+def dummyvacuums7(request):
+    request.cls.device = DummyVacuumS7()
+
+
+@pytest.mark.usefixtures("dummyvacuums7")
+class TestVacuumS7(TestCase):
+    def test_mop_intensity(self):
+        """Test getting mop intensity."""
+        with patch.object(self.device, "send", return_value=[203]) as mock_method:
+            assert self.device.mop_intensity()
+            mock_method.assert_called_once_with("get_water_box_custom_mode")
+
+    def test_set_mop_intensity(self):
+        """Test setting mop intensity."""
+        with patch.object(self.device, "send", return_value=[203]) as mock_method:
+            assert self.device.set_mop_intensity(MopIntensity.Intense)
+            mock_method.assert_called_once_with("set_water_box_custom_mode", [203])

miio/integrations/vacuum/roborock/vacuum.py

@@ -114,6 +114,15 @@ class MopMode(enum.Enum):
     Deep = 301
 
 
+class MopIntensity(enum.Enum):
+    """Mop scrub intensity on S7."""
+
+    Close = 200
+    Mild = 201
+    Moderate = 202
+    Intense = 203
+
+
 class CarpetCleaningMode(enum.Enum):
     """Type of carpet cleaning/avoidance."""
 
@@ -837,6 +846,22 @@ def set_mop_mode(self, mop_mode: MopMode):
         """Set mop mode setting."""
         return self.send("set_mop_mode", [mop_mode.value])[0] == "ok"
 
+    @command()
+    def mop_intensity(self) -> MopIntensity:
+        """Get mop scrub intensity setting."""
+        if self.model != ROCKROBO_S7:
+            raise VacuumException("Mop scrub intensity not supported by %s", self.model)
+
+        return MopIntensity(self.send("get_water_box_custom_mode")[0])
+
+    @command(click.argument("mop_intensity", type=EnumType(MopIntensity)))
+    def set_mop_intensity(self, mop_intensity: MopIntensity):
+        """Set mop scrub intensity setting."""
+        if self.model != ROCKROBO_S7:
+            raise VacuumException("Mop scrub intensity not supported by %s", self.model)
+
+        return self.send("set_water_box_custom_mode", [mop_intensity.value])
+
     @command()
     def child_lock(self) -> bool:
         """Get child lock setting."""