Add Xiaomi Miio vacuum config flow

homeassistant/components/xiaomi_miio/__init__.py

@@ -17,13 +17,15 @@
     DOMAIN,
     KEY_COORDINATOR,
     MODELS_SWITCH,
+    MODELS_VACUUM,
 )
 from .gateway import ConnectXiaomiGateway
 
 _LOGGER = logging.getLogger(__name__)
 
 GATEWAY_PLATFORMS = ["alarm_control_panel", "sensor", "light"]
 SWITCH_PLATFORMS = ["switch"]
+VACUUM_PLATFORMS = ["vacuum"]
 
 
 async def async_setup(hass: core.HomeAssistant, config: dict):
@@ -117,9 +119,14 @@ async def async_setup_device_entry(
     model = entry.data[CONF_MODEL]
 
     # Identify platforms to setup
+    platforms = []
     if model in MODELS_SWITCH:
         platforms = SWITCH_PLATFORMS
-    else:
+    for vacuum_model in MODELS_VACUUM:
+        if model.startswith(vacuum_model):
+            platforms = VACUUM_PLATFORMS
+
+    if not platforms:
         return False
 
     for component in platforms:

homeassistant/components/xiaomi_miio/config_flow.py

@@ -15,8 +15,9 @@
     CONF_MAC,
     CONF_MODEL,
     DOMAIN,
+    MODELS_ALL,
+    MODELS_ALL_DEVICES,
     MODELS_GATEWAY,
-    MODELS_SWITCH,
 )
 from .device import ConnectXiaomiDevice
 
@@ -29,6 +30,7 @@
     vol.Required(CONF_TOKEN): vol.All(str, vol.Length(min=32, max=32)),
 }
 DEVICE_CONFIG = vol.Schema({vol.Required(CONF_HOST): str}).extend(DEVICE_SETTINGS)
+DEVICE_MODEL_CONFIG = {vol.Optional(CONF_MODEL): vol.In(MODELS_ALL)}
 
 
 class XiaomiMiioFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
@@ -40,6 +42,7 @@ class XiaomiMiioFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
     def __init__(self):
         """Initialize."""
         self.host = None
+        self.mac = None
 
     async def async_step_import(self, conf: dict):
         """Import a configuration from config.yaml."""
@@ -53,15 +56,15 @@ async def async_step_zeroconf(self, discovery_info):
         """Handle zeroconf discovery."""
         name = discovery_info.get("name")
         self.host = discovery_info.get("host")
-        mac_address = discovery_info.get("properties", {}).get("mac")
+        self.mac = discovery_info.get("properties", {}).get("mac")
 
-        if not name or not self.host or not mac_address:
+        if not name or not self.host or not self.mac:
             return self.async_abort(reason="not_xiaomi_miio")
 
         # Check which device is discovered.
         for gateway_model in MODELS_GATEWAY:
             if name.startswith(gateway_model.replace(".", "-")):
-                unique_id = format_mac(mac_address)
+                unique_id = format_mac(self.mac)
                 await self.async_set_unique_id(unique_id)
                 self._abort_if_unique_id_configured({CONF_HOST: self.host})
 
@@ -70,9 +73,9 @@ async def async_step_zeroconf(self, discovery_info):
                 )
 
                 return await self.async_step_device()
-        for switch_model in MODELS_SWITCH:
-            if name.startswith(switch_model.replace(".", "-")):
-                unique_id = format_mac(mac_address)
+        for device_model in MODELS_ALL_DEVICES:
+            if name.startswith(device_model.replace(".", "-")):
+                unique_id = format_mac(self.mac)
                 await self.async_set_unique_id(unique_id)
                 self._abort_if_unique_id_configured({CONF_HOST: self.host})
 
@@ -95,6 +98,7 @@ async def async_step_device(self, user_input=None):
         errors = {}
         if user_input is not None:
             token = user_input[CONF_TOKEN]
+            model = user_input.get(CONF_MODEL)
             if user_input.get(CONF_HOST):
                 self.host = user_input[CONF_HOST]
 
@@ -103,12 +107,17 @@ async def async_step_device(self, user_input=None):
             await connect_device_class.async_connect_device(self.host, token)
             device_info = connect_device_class.device_info
 
-            if device_info is not None:
+            if model is None and device_info is not None:
+                model = device_info.model
+
+            if model is not None:
+                if self.mac is None and device_info is not None:
+                    self.mac = format_mac(device_info.mac_address)
+
                 # Setup Gateways
                 for gateway_model in MODELS_GATEWAY:
-                    if device_info.model.startswith(gateway_model):
-                        mac = format_mac(device_info.mac_address)
-                        unique_id = mac
+                    if model.startswith(gateway_model):
+                        unique_id = self.mac
                         await self.async_set_unique_id(unique_id)
                         self._abort_if_unique_id_configured()
                         return self.async_create_entry(
@@ -117,29 +126,29 @@ async def async_step_device(self, user_input=None):
                                 CONF_FLOW_TYPE: CONF_GATEWAY,
                                 CONF_HOST: self.host,
                                 CONF_TOKEN: token,
-                                CONF_MODEL: device_info.model,
-                                CONF_MAC: mac,
+                                CONF_MODEL: model,
+                                CONF_MAC: self.mac,
                             },
                         )
 
                 # Setup all other Miio Devices
                 name = user_input.get(CONF_NAME, DEFAULT_DEVICE_NAME)
 
-                if device_info.model in MODELS_SWITCH:
-                    mac = format_mac(device_info.mac_address)
-                    unique_id = mac
-                    await self.async_set_unique_id(unique_id)
-                    self._abort_if_unique_id_configured()
-                    return self.async_create_entry(
-                        title=name,
-                        data={
-                            CONF_FLOW_TYPE: CONF_DEVICE,
-                            CONF_HOST: self.host,
-                            CONF_TOKEN: token,
-                            CONF_MODEL: device_info.model,
-                            CONF_MAC: mac,
-                        },
-                    )
+                for device_model in MODELS_ALL_DEVICES:
+                    if model.startswith(device_model):
+                        unique_id = self.mac
+                        await self.async_set_unique_id(unique_id)
+                        self._abort_if_unique_id_configured()
+                        return self.async_create_entry(
+                            title=name,
+                            data={
+                                CONF_FLOW_TYPE: CONF_DEVICE,
+                                CONF_HOST: self.host,
+                                CONF_TOKEN: token,
+                                CONF_MODEL: model,
+                                CONF_MAC: self.mac,
+                            },
+                        )
                 errors["base"] = "unknown_device"
             else:
                 errors["base"] = "cannot_connect"
@@ -149,4 +158,7 @@ async def async_step_device(self, user_input=None):
         else:
             schema = DEVICE_CONFIG
 
+        if errors:
+            schema = schema.extend(DEVICE_MODEL_CONFIG)
+
         return self.async_show_form(step_id="device", data_schema=schema, errors=errors)

homeassistant/components/xiaomi_miio/const.py

@@ -23,6 +23,10 @@
     "chuangmi.plug.hmi206",
     "lumi.acpartner.v3",
 ]
+MODELS_VACUUM = ["roborock.vacuum"]
+
+MODELS_ALL_DEVICES = MODELS_SWITCH + MODELS_VACUUM
+MODELS_ALL = MODELS_ALL_DEVICES + MODELS_GATEWAY
 
 # Fan Services
 SERVICE_SET_BUZZER_ON = "fan_set_buzzer_on"

homeassistant/components/xiaomi_miio/device.py

@@ -78,10 +78,14 @@ def name(self):
     @property
     def device_info(self):
         """Return the device info."""
-        return {
-            "connections": {(dr.CONNECTION_NETWORK_MAC, self._mac)},
+        device_info = {
             "identifiers": {(DOMAIN, self._device_id)},
             "manufacturer": "Xiaomi",
             "name": self._name,
             "model": self._model,
         }
+
+        if self._mac is not None:
+            device_info["connections"] = {(dr.CONNECTION_NETWORK_MAC, self._mac)}
+
+        return device_info

homeassistant/components/xiaomi_miio/strings.json

@@ -8,7 +8,7 @@
         "data": {
           "host": "[%key:common::config_flow::data::ip%]",
           "token": "[%key:common::config_flow::data::api_token%]",
-          "name": "Name of the device"
+          "model": "Device model (Optional)"
         }
       }
     },

homeassistant/components/xiaomi_miio/translations/en.json

@@ -6,35 +6,18 @@
         },
         "error": {
             "cannot_connect": "Failed to connect",
-            "no_device_selected": "No device selected, please select one device.",
             "unknown_device": "The device model is not known, not able to setup the device using config flow."
         },
         "flow_title": "Xiaomi Miio: {name}",
         "step": {
             "device": {
                 "data": {
                     "host": "IP Address",
-                    "name": "Name of the device",
-                    "token": "API Token"
+                    "token": "API Token",
+                    "model": "Device model (Optional)"
                 },
                 "description": "You will need the 32 character API Token, see https://www.home-assistant.io/integrations/vacuum.xiaomi_miio/#retrieving-the-access-token for instructions. Please note, that this API Token is different from the key used by the Xiaomi Aqara integration.",
                 "title": "Connect to a Xiaomi Miio Device or Xiaomi Gateway"
-            },
-            "gateway": {
-                "data": {
-                    "host": "IP Address",
-                    "name": "Name of the Gateway",
-                    "token": "API Token"
-                },
-                "description": "You will need the 32 character API Token, see https://www.home-assistant.io/integrations/vacuum.xiaomi_miio/#retrieving-the-access-token for instructions. Please note, that this API Token is different from the key used by the Xiaomi Aqara integration.",
-                "title": "Connect to a Xiaomi Gateway"
-            },
-            "user": {
-                "data": {
-                    "gateway": "Connect to a Xiaomi Gateway"
-                },
-                "description": "Select to which device you want to connect.",
-                "title": "Xiaomi Miio"
             }
         }
     }