Add roborock mop features

homeassistant/components/xiaomi_miio/__init__.py

@@ -28,9 +28,11 @@
     FanZA5,
     RoborockVacuum,
     Timer,
+    VacuumException,
     VacuumStatus,
 )
 from miio.gateway.gateway import GatewayException
+from miio.integrations.vacuum.roborock.vacuum import MopIntensity, MopMode
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, CONF_MODEL, CONF_TOKEN, Platform
@@ -199,6 +201,8 @@ class VacuumCoordinatorData:
     timers: list[Timer]
     fan_speeds: dict[str, int]
     fan_speeds_reverse: dict[int, str]
+    mop_mode: MopMode | None
+    mop_intensity: MopIntensity | None
 
 
 @dataclass(init=False, frozen=True)
@@ -221,6 +225,8 @@ class VacuumCoordinatorDataAttributes:
     timer: str = "timer"
     fan_speeds: str = "fan_speeds"
     fan_speeds_reverse: str = "fan_speeds_reverse"
+    mop_mode: str = "mop_mode"
+    mop_intensity: str = "mop_intensity"
 
 
 def _async_update_data_vacuum(hass, device: RoborockVacuum):
@@ -237,6 +243,11 @@ def update() -> VacuumCoordinatorData:
             )
 
         fan_speeds = device.fan_speed_presets()
+        mop_mode = device.mop_mode()
+        try:
+            mop_intensity = device.mop_intensity()
+        except VacuumException:
+            mop_intensity = None
 
         data = VacuumCoordinatorData(
             device.status(),
@@ -247,6 +258,8 @@ def update() -> VacuumCoordinatorData:
             timer,
             fan_speeds,
             {v: k for k, v in fan_speeds.items()},
+            mop_mode,
+            mop_intensity,
         )
 
         return data

homeassistant/components/xiaomi_miio/const.py

@@ -299,6 +299,8 @@ class SetupException(Exception):
 SERVICE_SET_POWER_PRICE = "switch_set_power_price"
 
 # Vacuum Services
+SERVICE_SET_MOP_MODE = "set_mop_mode"
+SERVICE_SET_MOP_INTENSITY = "set_mop_intensity"
 SERVICE_MOVE_REMOTE_CONTROL = "vacuum_remote_control_move"
 SERVICE_MOVE_REMOTE_CONTROL_STEP = "vacuum_remote_control_move_step"
 SERVICE_START_REMOTE_CONTROL = "vacuum_remote_control_start"

homeassistant/components/xiaomi_miio/services.yaml

@@ -173,6 +173,32 @@ remote_set_led_off:
       integration: xiaomi_miio
       domain: remote
 
+set_mop_mode:
+  name: Set mop mode
+  description: "Set the mop mode of a vacuum."
+  target:
+    entity:
+      integration: xiaomi_miio
+      domain: vacuum
+  fields:
+    mop_mode:
+      name: Mop mode
+      description: "Define the mop mode to set in the vacuum."
+      required: true
+
+set_mop_intensity:
+  name: Set mop intensity
+  description: "Set the mop intensity of a vacuum."
+  target:
+    entity:
+      integration: xiaomi_miio
+      domain: vacuum
+  fields:
+    mop_intensity:
+      name: Mop intensity
+      description: "Define the mop intensity to set in the vacuum."
+      required: true
+
 switch_set_wifi_led_on:
   name: Switch set Wi-fi LED on
   description: Turn the wifi led on.

homeassistant/components/xiaomi_miio/vacuum.py

@@ -6,6 +6,7 @@
 from typing import Any
 
 from miio import DeviceException
+from miio.integrations.vacuum.roborock.vacuum import MopIntensity, MopMode
 import voluptuous as vol
 
 from homeassistant.components.vacuum import (
@@ -37,6 +38,8 @@
     SERVICE_GOTO,
     SERVICE_MOVE_REMOTE_CONTROL,
     SERVICE_MOVE_REMOTE_CONTROL_STEP,
+    SERVICE_SET_MOP_INTENSITY,
+    SERVICE_SET_MOP_MODE,
     SERVICE_START_REMOTE_CONTROL,
     SERVICE_STOP_REMOTE_CONTROL,
 )
@@ -45,6 +48,8 @@
 _LOGGER = logging.getLogger(__name__)
 
 ATTR_ERROR = "error"
+ATTR_MOP_MODE = "mop_mode"
+ATTR_MOP_INTENSITY = "mop_intensity"
 ATTR_RC_DURATION = "duration"
 ATTR_RC_ROTATION = "rotation"
 ATTR_RC_VELOCITY = "velocity"
@@ -177,6 +182,30 @@ async def async_setup_entry(
             {vol.Required("segments"): vol.Any(vol.Coerce(int), [vol.Coerce(int)])},
             MiroboVacuum.async_clean_segment.__name__,
         )
+        platform.async_register_entity_service(
+            SERVICE_SET_MOP_MODE,
+            {
+                vol.Required("mop_mode"): vol.Any(
+                    vol.Coerce(int),
+                    [vol.Coerce(int)],
+                    vol.Coerce(str),
+                    [vol.Coerce(str)],
+                )
+            },
+            MiroboVacuum.async_set_mop_mode.__name__,
+        )
+        platform.async_register_entity_service(
+            SERVICE_SET_MOP_INTENSITY,
+            {
+                vol.Required("mop_intensity"): vol.Any(
+                    vol.Coerce(int),
+                    [vol.Coerce(int)],
+                    vol.Coerce(str),
+                    [vol.Coerce(str)],
+                )
+            },
+            MiroboVacuum.async_set_mop_intensity.__name__,
+        )
 
     async_add_entities(entities, update_before_add=True)
 
@@ -249,6 +278,20 @@ def fan_speed_list(self) -> list[str]:
             return list(speed_list)
         return []
 
+    @property
+    def mop_mode(self) -> str | None:
+        """Return the mop mode of the vacuum cleaner."""
+        if mop_mode := self.coordinator.data.mop_mode:
+            return mop_mode.name
+        return None
+
+    @property
+    def mop_intensity(self) -> str | None:
+        """Return the mop intensity of the vacuum cleaner."""
+        if mop_intensity := self.coordinator.data.mop_intensity:
+            return mop_intensity.name
+        return None
+
     @property
     def timers(self) -> list[dict[str, Any]]:
         """Get the list of added timers of the vacuum cleaner."""
@@ -267,6 +310,12 @@ def extra_state_attributes(self) -> dict[str, Any]:
         attrs: dict[str, Any] = {}
         attrs[ATTR_STATUS] = str(self.coordinator.data.status.state)
 
+        if mop_mode := self.coordinator.data.mop_mode:
+            attrs[ATTR_MOP_MODE] = mop_mode.name
+
+        if mop_intensity := self.coordinator.data.mop_intensity:
+            attrs[ATTR_MOP_INTENSITY] = mop_intensity.name
+
         if self.coordinator.data.status.got_error:
             attrs[ATTR_ERROR] = self.coordinator.data.status.error
 
@@ -316,6 +365,42 @@ async def async_set_fan_speed(self, fan_speed: str, **kwargs: Any) -> None:
             "Unable to set fan speed: %s", self._device.set_fan_speed, fan_speed_int
         )
 
+    async def async_set_mop_mode(self, mop_mode: str, **kwargs: Any) -> None:
+        """Set mop mode."""
+        mop_mode_int = MopMode[mop_mode]
+        if mop_mode_int is None:
+            try:
+                mop_mode_int = int(mop_mode)
+            except ValueError as exc:
+                _LOGGER.error(
+                    "Mop mode not recognized (%s). Valid mop modes are: %s",
+                    exc,
+                    list(MopMode),
+                )
+                return
+        await self._try_command(
+            "Unable to set mop mode: %s", self._device.set_mop_mode, mop_mode_int
+        )
+
+    async def async_set_mop_intensity(self, mop_intensity: str, **kwargs: Any) -> None:
+        """Set mop intensity."""
+        mop_intensity_int = MopIntensity[mop_intensity]
+        if mop_intensity_int is None:
+            try:
+                mop_intensity_int = int(mop_intensity)
+            except ValueError as exc:
+                _LOGGER.error(
+                    "Mop intensity not recognized (%s). Valid mop intensities are: %s",
+                    exc,
+                    list(MopIntensity),
+                )
+                return
+        await self._try_command(
+            "Unable to set mop intensity: %s",
+            self._device.set_mop_intensity,
+            mop_intensity_int,
+        )
+
     async def async_return_to_base(self, **kwargs: Any) -> None:
         """Set the vacuum cleaner to return to the dock."""
         await self._try_command("Unable to return home: %s", self._device.home)