Add buttons for resetting vacuum consumable status in xiaomi_miio (#91483)

Co-authored-by: Franck Nijhof <frenck@frenck.nl>
Co-authored-by: Eugenio Panadero <eugenio.panadero@gmail.com>

homeassistant/components/xiaomi_miio/__init__.py

@@ -105,7 +105,12 @@ HUMIDIFIER_PLATFORMS = [
     Platform.SWITCH,
 ]
 LIGHT_PLATFORMS = [Platform.LIGHT]
-VACUUM_PLATFORMS = [Platform.BINARY_SENSOR, Platform.SENSOR, Platform.VACUUM]
+VACUUM_PLATFORMS = [
+    Platform.BINARY_SENSOR,
+    Platform.SENSOR,
+    Platform.BUTTON,
+    Platform.VACUUM,
+]
 AIR_MONITOR_PLATFORMS = [Platform.AIR_QUALITY, Platform.SENSOR]
 
 MODEL_TO_CLASS_MAP = {

homeassistant/components/xiaomi_miio/button.py

@@ -3,6 +3,8 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 
+from miio.integrations.vacuum.roborock.vacuum import Consumable
+
 from homeassistant.components.button import (
     ButtonDeviceClass,
     ButtonEntity,
@@ -19,22 +21,33 @@ from .const import (
     KEY_DEVICE,
     MODEL_AIRFRESH_A1,
     MODEL_AIRFRESH_T2017,
+    MODELS_VACUUM,
 )
 from .device import XiaomiCoordinatedMiioEntity
 
+# Fans
 ATTR_RESET_DUST_FILTER = "reset_dust_filter"
 ATTR_RESET_UPPER_FILTER = "reset_upper_filter"
 
+# Vacuums
+METHOD_VACUUM_RESET_CONSUMABLE = "consumable_reset"
+ATTR_RESET_VACUUM_MAIN_BRUSH = "reset_vacuum_main_brush"
+ATTR_RESET_VACUUM_SIDE_BRUSH = "reset_vacuum_side_brush"
+ATTR_RESET_VACUUM_FILTER = "reset_vacuum_filter"
+ATTR_RESET_VACUUM_SENSOR_DIRTY = "reset_vacuum_sensor_dirty"
+
 
 @dataclass
 class XiaomiMiioButtonDescription(ButtonEntityDescription):
     """A class that describes button entities."""
 
     method_press: str = ""
+    method_press_params: Consumable | None = None
     method_press_error_message: str = ""
 
 
 BUTTON_TYPES = (
+    # Fans
     XiaomiMiioButtonDescription(
         key=ATTR_RESET_DUST_FILTER,
         name="Reset dust filter",
@@ -51,6 +64,50 @@ BUTTON_TYPES = (
         method_press_error_message="Resetting the upper filter lifetime failed.",
         entity_category=EntityCategory.CONFIG,
     ),
+    # Vacuums
+    XiaomiMiioButtonDescription(
+        key=ATTR_RESET_VACUUM_MAIN_BRUSH,
+        name="Reset main brush",
+        icon="mdi:brush",
+        method_press=METHOD_VACUUM_RESET_CONSUMABLE,
+        method_press_params=Consumable.MainBrush,
+        method_press_error_message="Resetting the main brush lifetime failed.",
+        entity_category=EntityCategory.CONFIG,
+    ),
+    XiaomiMiioButtonDescription(
+        key=ATTR_RESET_VACUUM_SIDE_BRUSH,
+        name="Reset side brush",
+        icon="mdi:brush",
+        method_press=METHOD_VACUUM_RESET_CONSUMABLE,
+        method_press_params=Consumable.SideBrush,
+        method_press_error_message="Resetting the side brush lifetime failed.",
+        entity_category=EntityCategory.CONFIG,
+    ),
+    XiaomiMiioButtonDescription(
+        key=ATTR_RESET_VACUUM_FILTER,
+        name="Reset filter",
+        icon="mdi:air-filter",
+        method_press=METHOD_VACUUM_RESET_CONSUMABLE,
+        method_press_params=Consumable.Filter,
+        method_press_error_message="Resetting the filter lifetime failed.",
+        entity_category=EntityCategory.CONFIG,
+    ),
+    XiaomiMiioButtonDescription(
+        key=ATTR_RESET_VACUUM_SENSOR_DIRTY,
+        name="Reset sensor dirty",
+        icon="mdi:eye-outline",
+        method_press=METHOD_VACUUM_RESET_CONSUMABLE,
+        method_press_params=Consumable.SensorDirty,
+        method_press_error_message="Resetting the sensor lifetime failed.",
+        entity_category=EntityCategory.CONFIG,
+    ),
+)
+
+BUTTONS_FOR_VACUUM = (
+    ATTR_RESET_VACUUM_MAIN_BRUSH,
+    ATTR_RESET_VACUUM_SIDE_BRUSH,
+    ATTR_RESET_VACUUM_FILTER,
+    ATTR_RESET_VACUUM_SENSOR_DIRTY,
 )
 
 MODEL_TO_BUTTON_MAP: dict[str, tuple[str, ...]] = {
@@ -59,6 +116,7 @@ MODEL_TO_BUTTON_MAP: dict[str, tuple[str, ...]] = {
         ATTR_RESET_DUST_FILTER,
         ATTR_RESET_UPPER_FILTER,
     ),
+    **{model: BUTTONS_FOR_VACUUM for model in MODELS_VACUUM},
 }
 
 
@@ -114,4 +172,5 @@ class XiaomiGenericCoordinatedButton(XiaomiCoordinatedMiioEntity, ButtonEntity):
         await self._try_command(
             self.entity_description.method_press_error_message,
             method,
+            self.entity_description.method_press_params,
         )

tests/components/xiaomi_miio/test_button.py

@@ -0,0 +1,98 @@
+"""The tests for the xiaomi_miio button component."""
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from homeassistant.components.button import DOMAIN, SERVICE_PRESS
+from homeassistant.components.xiaomi_miio.const import (
+    CONF_DEVICE,
+    CONF_FLOW_TYPE,
+    CONF_MAC,
+    DOMAIN as XIAOMI_DOMAIN,
+    MODELS_VACUUM,
+)
+from homeassistant.const import (
+    ATTR_ENTITY_ID,
+    CONF_HOST,
+    CONF_MODEL,
+    CONF_TOKEN,
+    Platform,
+)
+from homeassistant.core import HomeAssistant
+from homeassistant.util import dt as dt_util
+
+from . import TEST_MAC
+
+from tests.common import MockConfigEntry
+
+
+@pytest.fixture(autouse=True)
+async def setup_test(hass: HomeAssistant):
+    """Initialize test xiaomi_miio for button entity."""
+
+    mock_vacuum = MagicMock()
+
+    with patch(
+        "homeassistant.components.xiaomi_miio.get_platforms",
+        return_value=[
+            Platform.BUTTON,
+        ],
+    ), patch("homeassistant.components.xiaomi_miio.RoborockVacuum") as mock_vacuum_cls:
+        mock_vacuum_cls.return_value = mock_vacuum
+        yield mock_vacuum
+
+
+async def test_vacuum_button_params(hass: HomeAssistant) -> None:
+    """Test the initial parameters of a vacuum button."""
+
+    entity_id = await setup_component(hass, "test_vacuum")
+
+    state = hass.states.get(f"{entity_id}_reset_main_brush")
+    assert state
+    assert state.state == "unknown"
+
+
+async def test_vacuum_button_press(hass: HomeAssistant) -> None:
+    """Test pressing a vacuum button."""
+
+    entity_id = await setup_component(hass, "test_vacuum")
+
+    state = hass.states.get(f"{entity_id}_reset_side_brush")
+    assert state
+    assert state.state == "unknown"
+
+    pressed_at = dt_util.utcnow()
+    await hass.services.async_call(
+        DOMAIN,
+        SERVICE_PRESS,
+        {ATTR_ENTITY_ID: entity_id + "_reset_side_brush"},
+        blocking=True,
+    )
+
+    state = hass.states.get(f"{entity_id}_reset_side_brush")
+    assert state
+    assert state.state[0:21] == pressed_at.isoformat()[0:21]  # drop millisecs
+
+
+async def setup_component(hass: HomeAssistant, entity_name: str) -> str:
+    """Set up vacuum component."""
+    entity_id = f"{DOMAIN}.{entity_name}"
+
+    config_entry = MockConfigEntry(
+        domain=XIAOMI_DOMAIN,
+        unique_id="123456",
+        title=entity_name,
+        data={
+            CONF_FLOW_TYPE: CONF_DEVICE,
+            CONF_HOST: "192.168.1.100",
+            CONF_TOKEN: "12345678901234567890123456789012",
+            CONF_MODEL: MODELS_VACUUM[0],
+            CONF_MAC: TEST_MAC,
+        },
+    )
+
+    config_entry.add_to_hass(hass)
+    await hass.config_entries.async_setup(config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    return entity_id