From 20b93132dd8b935995209bd8ba51b86b8baaa69e Mon Sep 17 00:00:00 2001
From: Kevin Hellemun <17928966+OGKevin@users.noreply.github.com>
Date: Mon, 8 Nov 2021 23:13:00 +0100
Subject: [PATCH] Support generic xiaomi_miio vacuums (#59317)

* Support generic xiaomi_miio vacuums

Signed-off-by: Kevin Hellemun <17928966+OGKevin@users.noreply.github.com>

* Fix lint

Signed-off-by: Kevin Hellemun <17928966+OGKevin@users.noreply.github.com>

* Remove warning log

Signed-off-by: Kevin Hellemun <17928966+OGKevin@users.noreply.github.com>
---
 homeassistant/components/xiaomi_miio/__init__.py | 12 ++++++++++--
 homeassistant/components/xiaomi_miio/const.py    |  4 +++-
 homeassistant/components/xiaomi_miio/sensor.py   | 10 ++++++++--
 3 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/homeassistant/components/xiaomi_miio/__init__.py b/homeassistant/components/xiaomi_miio/__init__.py
index 5513485f51e..9d983d32252 100644
--- a/homeassistant/components/xiaomi_miio/__init__.py
+++ b/homeassistant/components/xiaomi_miio/__init__.py
@@ -66,6 +66,8 @@ from .const import (
     MODELS_PURIFIER_MIOT,
     MODELS_SWITCH,
     MODELS_VACUUM,
+    ROBOROCK_GENERIC,
+    ROCKROBO_GENERIC,
     AuthException,
     SetupException,
 )
@@ -267,7 +269,7 @@ async def async_create_miio_device_and_coordinator(
     hass: core.HomeAssistant, entry: config_entries.ConfigEntry
 ):
     """Set up a data coordinator and one miio device to service multiple entities."""
-    model = entry.data[CONF_MODEL]
+    model: str = entry.data[CONF_MODEL]
     host = entry.data[CONF_HOST]
     token = entry.data[CONF_TOKEN]
     name = entry.title
@@ -280,6 +282,8 @@ async def async_create_miio_device_and_coordinator(
         model not in MODELS_HUMIDIFIER
         and model not in MODELS_FAN
         and model not in MODELS_VACUUM
+        and not model.startswith(ROBOROCK_GENERIC)
+        and not model.startswith(ROCKROBO_GENERIC)
     ):
         return
 
@@ -304,7 +308,11 @@ async def async_create_miio_device_and_coordinator(
         device = AirPurifier(host, token)
     elif model.startswith("zhimi.airfresh."):
         device = AirFresh(host, token)
-    elif model in MODELS_VACUUM:
+    elif (
+        model in MODELS_VACUUM
+        or model.startswith(ROBOROCK_GENERIC)
+        or model.startswith(ROCKROBO_GENERIC)
+    ):
         device = Vacuum(host, token)
         update_method = _async_update_data_vacuum
         coordinator_class = DataUpdateCoordinator[VacuumCoordinatorData]
diff --git a/homeassistant/components/xiaomi_miio/const.py b/homeassistant/components/xiaomi_miio/const.py
index 60c16a0717a..8c83c8015b2 100644
--- a/homeassistant/components/xiaomi_miio/const.py
+++ b/homeassistant/components/xiaomi_miio/const.py
@@ -202,7 +202,8 @@ ROCKROBO_S4_MAX = "roborock.vacuum.a19"
 ROCKROBO_S5_MAX = "roborock.vacuum.s5e"
 ROCKROBO_S6_PURE = "roborock.vacuum.a08"
 ROCKROBO_E2 = "roborock.vacuum.e2"
-ROCKROBO_GENERIC = "roborock.vacuum"
+ROBOROCK_GENERIC = "roborock.vacuum"
+ROCKROBO_GENERIC = "rockrobo.vacuum"
 MODELS_VACUUM = [
     ROCKROBO_V1,
     ROCKROBO_E2,
@@ -214,6 +215,7 @@ MODELS_VACUUM = [
     ROCKROBO_S6_MAXV,
     ROCKROBO_S6_PURE,
     ROCKROBO_S7,
+    ROBOROCK_GENERIC,
     ROCKROBO_GENERIC,
 ]
 MODELS_VACUUM_WITH_MOP = [
diff --git a/homeassistant/components/xiaomi_miio/sensor.py b/homeassistant/components/xiaomi_miio/sensor.py
index ac26bc97fce..0d67014ced9 100644
--- a/homeassistant/components/xiaomi_miio/sensor.py
+++ b/homeassistant/components/xiaomi_miio/sensor.py
@@ -81,6 +81,8 @@ from .const import (
     MODELS_PURIFIER_MIIO,
     MODELS_PURIFIER_MIOT,
     MODELS_VACUUM,
+    ROBOROCK_GENERIC,
+    ROCKROBO_GENERIC,
 )
 from .device import XiaomiCoordinatedMiioEntity, XiaomiMiioEntity
 from .gateway import XiaomiGatewayDevice
@@ -592,7 +594,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
     elif config_entry.data[CONF_FLOW_TYPE] == CONF_DEVICE:
         host = config_entry.data[CONF_HOST]
         token = config_entry.data[CONF_TOKEN]
-        model = config_entry.data[CONF_MODEL]
+        model: str = config_entry.data[CONF_MODEL]
 
         if model in (MODEL_FAN_ZA1, MODEL_FAN_ZA3, MODEL_FAN_ZA4, MODEL_FAN_P5):
             return
@@ -624,7 +626,11 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
                 sensors = PURIFIER_MIIO_SENSORS
             elif model in MODELS_PURIFIER_MIOT:
                 sensors = PURIFIER_MIOT_SENSORS
-            elif model in MODELS_VACUUM:
+            elif (
+                model in MODELS_VACUUM
+                or model.startswith(ROBOROCK_GENERIC)
+                or model.startswith(ROCKROBO_GENERIC)
+            ):
                 return _setup_vacuum_sensors(hass, config_entry, async_add_entities)
 
             for sensor, description in SENSOR_TYPES.items():