Improve Xioami Aqara zeroconf discovery handling (#37469)

Co-authored-by: Franck Nijhof <git@frenck.dev>
Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

homeassistant/components/xiaomi_aqara/__init__.py

@@ -142,11 +142,11 @@ async def async_setup_entry(
     xiaomi_gateway = await hass.async_add_executor_job(
         XiaomiGateway,
         entry.data[CONF_HOST],
-        entry.data[CONF_PORT],
         entry.data[CONF_SID],
         entry.data[CONF_KEY],
         DEFAULT_DISCOVERY_RETRY,
         entry.data[CONF_INTERFACE],
+        entry.data[CONF_PORT],
         entry.data[CONF_PROTOCOL],
     )
     hass.data[DOMAIN][GATEWAYS_KEY][entry.entry_id] = xiaomi_gateway

homeassistant/components/xiaomi_aqara/config_flow.py

@@ -3,10 +3,11 @@ import logging
 from socket import gaierror
 
 import voluptuous as vol
-from xiaomi_gateway import XiaomiGatewayDiscovery
+from xiaomi_gateway import MULTICAST_PORT, XiaomiGateway, XiaomiGatewayDiscovery
 
 from homeassistant import config_entries
 from homeassistant.const import CONF_HOST, CONF_MAC, CONF_NAME, CONF_PORT
+from homeassistant.core import callback
 from homeassistant.helpers.device_registry import format_mac
 
 # pylint: disable=unused-import
@@ -15,6 +16,7 @@ from .const import (
     CONF_KEY,
     CONF_PROTOCOL,
     CONF_SID,
+    DEFAULT_DISCOVERY_RETRY,
     DOMAIN,
     ZEROCONF_GATEWAY,
 )
@@ -28,6 +30,11 @@ DEFAULT_INTERFACE = "any"
 GATEWAY_CONFIG = vol.Schema(
     {vol.Optional(CONF_INTERFACE, default=DEFAULT_INTERFACE): str}
 )
+CONFIG_HOST = {
+    vol.Optional(CONF_HOST): str,
+    vol.Optional(CONF_MAC): str,
+}
+GATEWAY_CONFIG_HOST = GATEWAY_CONFIG.extend(CONFIG_HOST)
 GATEWAY_SETTINGS = vol.Schema(
     {
         vol.Optional(CONF_KEY): vol.All(str, vol.Length(min=16, max=16)),
@@ -46,44 +53,78 @@ class XiaomiAqaraFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
         """Initialize."""
         self.host = None
         self.interface = DEFAULT_INTERFACE
+        self.sid = None
         self.gateways = None
         self.selected_gateway = None
 
+    @callback
+    def async_show_form_step_user(self, errors):
+        """Show the form belonging to the user step."""
+        schema = GATEWAY_CONFIG
+        if (self.host is None and self.sid is None) or errors:
+            schema = GATEWAY_CONFIG_HOST
+
+        return self.async_show_form(step_id="user", data_schema=schema, errors=errors)
+
     async def async_step_user(self, user_input=None):
         """Handle a flow initialized by the user."""
         errors = {}
-        if user_input is not None:
+        if user_input is None:
-            self.interface = user_input[CONF_INTERFACE]
+            return self.async_show_form_step_user(errors)
 
-            # Discover Xiaomi Aqara Gateways in the netwerk to get required SIDs.
+        self.interface = user_input[CONF_INTERFACE]
-            xiaomi = XiaomiGatewayDiscovery(self.hass.add_job, [], self.interface)
-            try:
-                await self.hass.async_add_executor_job(xiaomi.discover_gateways)
-            except gaierror:
-                errors[CONF_INTERFACE] = "invalid_interface"
 
-            if not errors:
+        # allow optional manual setting of host and mac
-                self.gateways = xiaomi.gateways
+        if self.host is None and self.sid is None:
+            self.host = user_input.get(CONF_HOST)
+            mac_address = user_input.get(CONF_MAC)
 
-                # if host is already known by zeroconf discovery
+            # format sid from mac_address
-                if self.host is not None:
+            if mac_address is not None:
-                    self.selected_gateway = self.gateways.get(self.host)
+                self.sid = format_mac(mac_address).replace(":", "")
-                    if self.selected_gateway is not None:
-                        return await self.async_step_settings()
 
-                    errors["base"] = "not_found_error"
+        # if host is already known by zeroconf discovery or manual optional settings
-                else:
+        if self.host is not None and self.sid is not None:
-                    if len(self.gateways) == 1:
+            # Connect to Xiaomi Aqara Gateway
-                        self.selected_gateway = list(self.gateways.values())[0]
+            self.selected_gateway = await self.hass.async_add_executor_job(
-                        return await self.async_step_settings()
+                XiaomiGateway,
-                    if len(self.gateways) > 1:
+                self.host,
-                        return await self.async_step_select()
+                self.sid,
+                None,
+                DEFAULT_DISCOVERY_RETRY,
+                self.interface,
+                MULTICAST_PORT,
+                None,
+            )
 
-                    errors["base"] = "discovery_error"
+            if self.selected_gateway.connection_error:
+                errors[CONF_HOST] = "invalid_host"
+            if self.selected_gateway.mac_error:
+                errors[CONF_MAC] = "invalid_mac"
+            if errors:
+                return self.async_show_form_step_user(errors)
 
-        return self.async_show_form(
+            return await self.async_step_settings()
-            step_id="user", data_schema=GATEWAY_CONFIG, errors=errors
+
-        )
+        # Discover Xiaomi Aqara Gateways in the netwerk to get required SIDs.
+        xiaomi = XiaomiGatewayDiscovery(self.hass.add_job, [], self.interface)
+        try:
+            await self.hass.async_add_executor_job(xiaomi.discover_gateways)
+        except gaierror:
+            errors[CONF_INTERFACE] = "invalid_interface"
+            return self.async_show_form_step_user(errors)
+
+        self.gateways = xiaomi.gateways
+
+        if len(self.gateways) == 1:
+            self.selected_gateway = list(self.gateways.values())[0]
+            self.sid = self.selected_gateway.sid
+            return await self.async_step_settings()
+        if len(self.gateways) > 1:
+            return await self.async_step_select()
+
+        errors["base"] = "discovery_error"
+        return self.async_show_form_step_user(errors)
 
     async def async_step_select(self, user_input=None):
         """Handle multiple aqara gateways found."""
@@ -91,6 +132,7 @@ class XiaomiAqaraFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
         if user_input is not None:
             ip_adress = user_input["select_ip"]
             self.selected_gateway = self.gateways[ip_adress]
+            self.sid = self.selected_gateway.sid
             return await self.async_step_settings()
 
         select_schema = vol.Schema(
@@ -123,9 +165,12 @@ class XiaomiAqaraFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
             )
             return self.async_abort(reason="not_xiaomi_aqara")
 
-        # format mac (include semicolns and make uppercase)
+        # format mac (include semicolns and make lowercase)
         mac_address = format_mac(mac_address)
 
+        # format sid from mac_address
+        self.sid = mac_address.replace(":", "")
+
         unique_id = mac_address
         await self.async_set_unique_id(unique_id)
         self._abort_if_unique_id_configured({CONF_HOST: self.host})
@@ -144,19 +189,18 @@ class XiaomiAqaraFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
             key = user_input.get(CONF_KEY)
             ip_adress = self.selected_gateway.ip_adress
             port = self.selected_gateway.port
-            sid = self.selected_gateway.sid
             protocol = self.selected_gateway.proto
 
             if key is not None:
                 # validate key by issuing stop ringtone playback command.
                 self.selected_gateway.key = key
-                valid_key = self.selected_gateway.write_to_hub(sid, mid=10000)
+                valid_key = self.selected_gateway.write_to_hub(self.sid, mid=10000)
             else:
                 valid_key = True
 
             if valid_key:
                 # format_mac, for a gateway the sid equels the mac address
-                mac_address = format_mac(sid)
+                mac_address = format_mac(self.sid)
 
                 # set unique_id
                 unique_id = mac_address
@@ -172,7 +216,7 @@ class XiaomiAqaraFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
                         CONF_INTERFACE: self.interface,
                         CONF_PROTOCOL: protocol,
                         CONF_KEY: key,
-                        CONF_SID: sid,
+                        CONF_SID: self.sid,
                     },
                 )
 

homeassistant/components/xiaomi_aqara/strings.json

@@ -4,9 +4,11 @@
     "step": {
       "user": {
         "title": "Xiaomi Aqara Gateway",
-        "description": "Connect to your Xiaomi Aqara Gateway",
+        "description": "Connect to your Xiaomi Aqara Gateway, if the IP and mac addresses are left empty, auto-discovery is used",
         "data": {
-          "interface": "The network interface to use"
+          "interface": "The network interface to use",
+          "host": "[%key:common::config_flow::data::ip%] (optional)",
+          "mac": "Mac Address (optional)"
         }
       },
       "settings": {
@@ -27,9 +29,10 @@
     },
     "error": {
       "discovery_error": "Failed to discover a Xiaomi Aqara Gateway, try using the IP of the device running HomeAssistant as interface",
-      "not_found_error": "Zeroconf discovered Gateway could not be located to get the necessary information, try using the IP of the device running HomeAssistant as interface",
       "invalid_interface": "Invalid network interface",
-      "invalid_key": "Invalid gateway key"
+      "invalid_key": "Invalid gateway key",
+      "invalid_host": "Invalid [%key:common::config_flow::data::ip%]",
+      "invalid_mac": "Invalid Mac Address"
     },
     "abort": {
       "already_configured": "[%key:common::config_flow::abort::already_configured_device%]",

homeassistant/components/xiaomi_aqara/translations/en.json

@@ -1,15 +1,16 @@
 {
     "config": {
         "abort": {
-            "already_configured": "Device is already configured",
+            "already_configured": "[%key:common::config_flow::abort::already_configured_device%]",
             "already_in_progress": "Config flow for this gateway is already in progress",
             "not_xiaomi_aqara": "Not a Xiaomi Aqara Gateway, discovered device did not match known gateways"
         },
         "error": {
             "discovery_error": "Failed to discover a Xiaomi Aqara Gateway, try using the IP of the device running HomeAssistant as interface",
+            "invalid_host": "Invalid [%key:common::config_flow::data::ip%]",
             "invalid_interface": "Invalid network interface",
             "invalid_key": "Invalid gateway key",
-            "not_found_error": "Zeroconf discovered Gateway could not be located to get the necessary information, try using the IP of the device running HomeAssistant as interface"
+            "invalid_mac": "Invalid Mac Address"
         },
         "flow_title": "Xiaomi Aqara Gateway: {name}",
         "step": {
@@ -30,9 +31,11 @@
             },
             "user": {
                 "data": {
-                    "interface": "The network interface to use"
+                    "host": "[%key:common::config_flow::data::ip%] (optional)",
+                    "interface": "The network interface to use",
+                    "mac": "Mac Address (optional)"
                 },
-                "description": "Connect to your Xiaomi Aqara Gateway",
+                "description": "Connect to your Xiaomi Aqara Gateway, if the IP and mac addresses are left empty, auto-discovery is used",
                 "title": "Xiaomi Aqara Gateway"
             }
         }

tests/components/xiaomi_aqara/test_config_flow.py

@@ -34,13 +34,22 @@ def xiaomi_aqara_fixture():
     with patch(
         "homeassistant.components.xiaomi_aqara.config_flow.XiaomiGatewayDiscovery",
         return_value=mock_gateway_discovery,
+    ), patch(
+        "homeassistant.components.xiaomi_aqara.config_flow.XiaomiGateway",
+        return_value=mock_gateway_discovery.gateways[TEST_HOST],
     ), patch(
         "homeassistant.components.xiaomi_aqara.async_setup_entry", return_value=True
     ):
         yield
 
 
-def get_mock_discovery(host_list, invalid_interface=False, invalid_key=False):
+def get_mock_discovery(
+    host_list,
+    invalid_interface=False,
+    invalid_key=False,
+    invalid_host=False,
+    invalid_mac=False,
+):
     """Return a mock gateway info instance."""
     gateway_discovery = Mock()
 
@@ -52,6 +61,8 @@ def get_mock_discovery(host_list, invalid_interface=False, invalid_key=False):
         gateway.port = TEST_PORT
         gateway.sid = TEST_SID
         gateway.proto = TEST_PROTOCOL
+        gateway.connection_error = invalid_host
+        gateway.mac_error = invalid_mac
 
         if invalid_key:
             gateway.write_to_hub = Mock(return_value=False)
@@ -185,6 +196,52 @@ async def test_config_flow_user_no_key_success(hass):
     }
 
 
+async def test_config_flow_user_host_mac_success(hass):
+    """Test a successful config flow initialized by the user with a host and mac specified."""
+    result = await hass.config_entries.flow.async_init(
+        const.DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+
+    assert result["type"] == "form"
+    assert result["step_id"] == "user"
+    assert result["errors"] == {}
+
+    mock_gateway_discovery = get_mock_discovery([])
+
+    with patch(
+        "homeassistant.components.xiaomi_aqara.config_flow.XiaomiGatewayDiscovery",
+        return_value=mock_gateway_discovery,
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {
+                const.CONF_INTERFACE: config_flow.DEFAULT_INTERFACE,
+                CONF_HOST: TEST_HOST,
+                CONF_MAC: TEST_MAC,
+            },
+        )
+
+    assert result["type"] == "form"
+    assert result["step_id"] == "settings"
+    assert result["errors"] == {}
+
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"], {CONF_NAME: TEST_NAME},
+    )
+
+    assert result["type"] == "create_entry"
+    assert result["title"] == TEST_NAME
+    assert result["data"] == {
+        CONF_HOST: TEST_HOST,
+        CONF_PORT: TEST_PORT,
+        CONF_MAC: TEST_MAC,
+        const.CONF_INTERFACE: config_flow.DEFAULT_INTERFACE,
+        const.CONF_PROTOCOL: TEST_PROTOCOL,
+        const.CONF_KEY: None,
+        const.CONF_SID: TEST_SID,
+    }
+
+
 async def test_config_flow_user_discovery_error(hass):
     """Test a failed config flow initialized by the user with no gateways discoverd."""
     result = await hass.config_entries.flow.async_init(
@@ -235,6 +292,66 @@ async def test_config_flow_user_invalid_interface(hass):
     assert result["errors"] == {const.CONF_INTERFACE: "invalid_interface"}
 
 
+async def test_config_flow_user_invalid_host(hass):
+    """Test a failed config flow initialized by the user with an invalid host."""
+    result = await hass.config_entries.flow.async_init(
+        const.DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+
+    assert result["type"] == "form"
+    assert result["step_id"] == "user"
+    assert result["errors"] == {}
+
+    mock_gateway_discovery = get_mock_discovery([TEST_HOST], invalid_host=True)
+
+    with patch(
+        "homeassistant.components.xiaomi_aqara.config_flow.XiaomiGateway",
+        return_value=mock_gateway_discovery.gateways[TEST_HOST],
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {
+                const.CONF_INTERFACE: config_flow.DEFAULT_INTERFACE,
+                CONF_HOST: "0.0.0.0",
+                CONF_MAC: TEST_MAC,
+            },
+        )
+
+    assert result["type"] == "form"
+    assert result["step_id"] == "user"
+    assert result["errors"] == {"host": "invalid_host"}
+
+
+async def test_config_flow_user_invalid_mac(hass):
+    """Test a failed config flow initialized by the user with an invalid mac."""
+    result = await hass.config_entries.flow.async_init(
+        const.DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+
+    assert result["type"] == "form"
+    assert result["step_id"] == "user"
+    assert result["errors"] == {}
+
+    mock_gateway_discovery = get_mock_discovery([TEST_HOST], invalid_mac=True)
+
+    with patch(
+        "homeassistant.components.xiaomi_aqara.config_flow.XiaomiGateway",
+        return_value=mock_gateway_discovery.gateways[TEST_HOST],
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {
+                const.CONF_INTERFACE: config_flow.DEFAULT_INTERFACE,
+                CONF_HOST: TEST_HOST,
+                CONF_MAC: "in:va:li:d0:0m:ac",
+            },
+        )
+
+    assert result["type"] == "form"
+    assert result["step_id"] == "user"
+    assert result["errors"] == {"mac": "invalid_mac"}
+
+
 async def test_config_flow_user_invalid_key(hass):
     """Test a failed config flow initialized by the user with an invalid key."""
     result = await hass.config_entries.flow.async_init(
@@ -335,34 +452,3 @@ async def test_zeroconf_unknown_device(hass):
 
     assert result["type"] == "abort"
     assert result["reason"] == "not_xiaomi_aqara"
-
-
-async def test_zeroconf_not_found_error(hass):
-    """Test a failed zeroconf discovery because the correct gateway could not be found."""
-    result = await hass.config_entries.flow.async_init(
-        const.DOMAIN,
-        context={"source": config_entries.SOURCE_ZEROCONF},
-        data={
-            zeroconf.ATTR_HOST: TEST_HOST,
-            ZEROCONF_NAME: TEST_ZEROCONF_NAME,
-            ZEROCONF_PROP: {ZEROCONF_MAC: TEST_MAC},
-        },
-    )
-
-    assert result["type"] == "form"
-    assert result["step_id"] == "user"
-    assert result["errors"] == {}
-
-    mock_gateway_discovery = get_mock_discovery([TEST_HOST_2])
-
-    with patch(
-        "homeassistant.components.xiaomi_aqara.config_flow.XiaomiGatewayDiscovery",
-        return_value=mock_gateway_discovery,
-    ):
-        result = await hass.config_entries.flow.async_configure(
-            result["flow_id"], {const.CONF_INTERFACE: config_flow.DEFAULT_INTERFACE},
-        )
-
-    assert result["type"] == "form"
-    assert result["step_id"] == "user"
-    assert result["errors"] == {"base": "not_found_error"}