From aaad9860021b45a08c663162cfcd38d414018f94 Mon Sep 17 00:00:00 2001 From: starkillerOG Date: Thu, 6 Aug 2020 12:54:18 +0200 Subject: [PATCH] Improve Xioami Aqara zeroconf discovery handling (#37469) Co-authored-by: Franck Nijhof Co-authored-by: Martin Hjelmare --- .../components/xiaomi_aqara/__init__.py | 2 +- .../components/xiaomi_aqara/config_flow.py | 108 +++++++++---- .../components/xiaomi_aqara/strings.json | 11 +- .../xiaomi_aqara/translations/en.json | 11 +- .../xiaomi_aqara/test_config_flow.py | 150 ++++++++++++++---- 5 files changed, 209 insertions(+), 73 deletions(-) diff --git a/homeassistant/components/xiaomi_aqara/__init__.py b/homeassistant/components/xiaomi_aqara/__init__.py index d759785f49f..c5b74e68af5 100644 --- a/homeassistant/components/xiaomi_aqara/__init__.py +++ b/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 diff --git a/homeassistant/components/xiaomi_aqara/config_flow.py b/homeassistant/components/xiaomi_aqara/config_flow.py index b9cfe58ac4b..fb66be76635 100644 --- a/homeassistant/components/xiaomi_aqara/config_flow.py +++ b/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: - self.interface = user_input[CONF_INTERFACE] + if user_input is None: + return self.async_show_form_step_user(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" + self.interface = user_input[CONF_INTERFACE] - if not errors: - self.gateways = xiaomi.gateways + # allow optional manual setting of host and mac + 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 - if self.host is not None: - self.selected_gateway = self.gateways.get(self.host) - if self.selected_gateway is not None: - return await self.async_step_settings() + # format sid from mac_address + if mac_address is not None: + self.sid = format_mac(mac_address).replace(":", "") - errors["base"] = "not_found_error" - else: - if len(self.gateways) == 1: - self.selected_gateway = list(self.gateways.values())[0] - return await self.async_step_settings() - if len(self.gateways) > 1: - return await self.async_step_select() + # if host is already known by zeroconf discovery or manual optional settings + if self.host is not None and self.sid is not None: + # Connect to Xiaomi Aqara Gateway + self.selected_gateway = await self.hass.async_add_executor_job( + XiaomiGateway, + self.host, + 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( - step_id="user", data_schema=GATEWAY_CONFIG, errors=errors - ) + return await self.async_step_settings() + + # 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, }, ) diff --git a/homeassistant/components/xiaomi_aqara/strings.json b/homeassistant/components/xiaomi_aqara/strings.json index 87e1d37cb93..5cbdc91a661 100644 --- a/homeassistant/components/xiaomi_aqara/strings.json +++ b/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%]", diff --git a/homeassistant/components/xiaomi_aqara/translations/en.json b/homeassistant/components/xiaomi_aqara/translations/en.json index 7b801e33089..b9f6fa7ab2a 100644 --- a/homeassistant/components/xiaomi_aqara/translations/en.json +++ b/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" } } diff --git a/tests/components/xiaomi_aqara/test_config_flow.py b/tests/components/xiaomi_aqara/test_config_flow.py index b7762317fdf..06fda84c934 100644 --- a/tests/components/xiaomi_aqara/test_config_flow.py +++ b/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"}