From c8451001a00856169efd3e92c06cfd83a2c7ef2b Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Mon, 22 Nov 2021 15:19:01 +0100
Subject: [PATCH] Use ZeroconfServiceInfo in xiaomi_aqara (#60131)

Co-authored-by: epenet <epenet@users.noreply.github.com>
---
 .../components/xiaomi_aqara/config_flow.py    | 12 ++++++---
 .../xiaomi_aqara/test_config_flow.py          | 25 +++++++++++--------
 2 files changed, 22 insertions(+), 15 deletions(-)

diff --git a/homeassistant/components/xiaomi_aqara/config_flow.py b/homeassistant/components/xiaomi_aqara/config_flow.py
index 3fbe7a71496..2b486cb87ad 100644
--- a/homeassistant/components/xiaomi_aqara/config_flow.py
+++ b/homeassistant/components/xiaomi_aqara/config_flow.py
@@ -6,8 +6,10 @@ import voluptuous as vol
 from xiaomi_gateway import MULTICAST_PORT, XiaomiGateway, XiaomiGatewayDiscovery
 
 from homeassistant import config_entries
+from homeassistant.components import zeroconf
 from homeassistant.const import CONF_HOST, CONF_MAC, CONF_NAME, CONF_PORT, CONF_PROTOCOL
 from homeassistant.core import callback
+from homeassistant.data_entry_flow import FlowResult
 from homeassistant.helpers.device_registry import format_mac
 
 from .const import (
@@ -144,11 +146,13 @@ class XiaomiAqaraFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
             step_id="select", data_schema=select_schema, errors=errors
         )
 
-    async def async_step_zeroconf(self, discovery_info):
+    async def async_step_zeroconf(
+        self, discovery_info: zeroconf.ZeroconfServiceInfo
+    ) -> FlowResult:
         """Handle zeroconf discovery."""
-        name = discovery_info.get("name")
-        self.host = discovery_info.get("host")
-        mac_address = discovery_info.get("properties", {}).get("mac")
+        name = discovery_info[zeroconf.ATTR_NAME]
+        self.host = discovery_info[zeroconf.ATTR_HOST]
+        mac_address = discovery_info[zeroconf.ATTR_PROPERTIES].get("mac")
 
         if not name or not self.host or not mac_address:
             return self.async_abort(reason="not_xiaomi_aqara")
diff --git a/tests/components/xiaomi_aqara/test_config_flow.py b/tests/components/xiaomi_aqara/test_config_flow.py
index 3f445a1fdec..27000af8a05 100644
--- a/tests/components/xiaomi_aqara/test_config_flow.py
+++ b/tests/components/xiaomi_aqara/test_config_flow.py
@@ -5,6 +5,7 @@ from unittest.mock import Mock, patch
 import pytest
 
 from homeassistant import config_entries
+from homeassistant.components import zeroconf
 from homeassistant.components.xiaomi_aqara import config_flow, const
 from homeassistant.const import CONF_HOST, CONF_MAC, CONF_NAME, CONF_PORT, CONF_PROTOCOL
 
@@ -400,11 +401,11 @@ async def test_zeroconf_success(hass):
     result = await hass.config_entries.flow.async_init(
         const.DOMAIN,
         context={"source": config_entries.SOURCE_ZEROCONF},
-        data={
-            CONF_HOST: TEST_HOST,
-            ZEROCONF_NAME: TEST_ZEROCONF_NAME,
-            ZEROCONF_PROP: {ZEROCONF_MAC: TEST_MAC},
-        },
+        data=zeroconf.ZeroconfServiceInfo(
+            host=TEST_HOST,
+            name=TEST_ZEROCONF_NAME,
+            properties={ZEROCONF_MAC: TEST_MAC},
+        ),
     )
 
     assert result["type"] == "form"
@@ -443,7 +444,9 @@ async def test_zeroconf_missing_data(hass):
     result = await hass.config_entries.flow.async_init(
         const.DOMAIN,
         context={"source": config_entries.SOURCE_ZEROCONF},
-        data={CONF_HOST: TEST_HOST, ZEROCONF_NAME: TEST_ZEROCONF_NAME},
+        data=zeroconf.ZeroconfServiceInfo(
+            host=TEST_HOST, name=TEST_ZEROCONF_NAME, properties={}
+        ),
     )
 
     assert result["type"] == "abort"
@@ -455,11 +458,11 @@ async def test_zeroconf_unknown_device(hass):
     result = await hass.config_entries.flow.async_init(
         const.DOMAIN,
         context={"source": config_entries.SOURCE_ZEROCONF},
-        data={
-            CONF_HOST: TEST_HOST,
-            ZEROCONF_NAME: "not-a-xiaomi-aqara-gateway",
-            ZEROCONF_PROP: {ZEROCONF_MAC: TEST_MAC},
-        },
+        data=zeroconf.ZeroconfServiceInfo(
+            host=TEST_HOST,
+            name="not-a-xiaomi-aqara-gateway",
+            properties={ZEROCONF_MAC: TEST_MAC},
+        ),
     )
 
     assert result["type"] == "abort"