Use ServiceInfo in yeelight (#60127)

Co-authored-by: epenet <epenet@users.noreply.github.com>

homeassistant/components/yeelight/config_flow.py

@@ -8,10 +8,11 @@ from yeelight.aio import AsyncBulb
 from yeelight.main import get_known_models
 
 from homeassistant import config_entries, exceptions
-from homeassistant.components.dhcp import IP_ADDRESS
+from homeassistant.components import dhcp, zeroconf
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.const import CONF_DEVICE, CONF_HOST, CONF_ID, CONF_NAME
 from homeassistant.core import callback
+from homeassistant.data_entry_flow import FlowResult
 import homeassistant.helpers.config_validation as cv
 
 from . import (
@@ -53,21 +54,25 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         self._discovered_model = None
         self._discovered_ip = None
 
-    async def async_step_homekit(self, discovery_info):
+    async def async_step_homekit(
+        self, discovery_info: zeroconf.ZeroconfServiceInfo
+    ) -> FlowResult:
         """Handle discovery from homekit."""
-        self._discovered_ip = discovery_info["host"]
+        self._discovered_ip = discovery_info[zeroconf.ATTR_HOST]
         return await self._async_handle_discovery()
 
-    async def async_step_dhcp(self, discovery_info):
+    async def async_step_dhcp(self, discovery_info: dhcp.DhcpServiceInfo) -> FlowResult:
         """Handle discovery from dhcp."""
-        self._discovered_ip = discovery_info[IP_ADDRESS]
+        self._discovered_ip = discovery_info[dhcp.IP_ADDRESS]
         return await self._async_handle_discovery()
 
-    async def async_step_zeroconf(self, discovery_info):
+    async def async_step_zeroconf(
+        self, discovery_info: zeroconf.ZeroconfServiceInfo
+    ) -> FlowResult:
         """Handle discovery from zeroconf."""
-        self._discovered_ip = discovery_info["host"]
+        self._discovered_ip = discovery_info[zeroconf.ATTR_HOST]
         await self.async_set_unique_id(
-            "{0:#0{1}x}".format(int(discovery_info["name"][-26:-18]), 18)
+            "{0:#0{1}x}".format(int(discovery_info[zeroconf.ATTR_NAME][-26:-18]), 18)
         )
         return await self._async_handle_discovery_with_unique_id()
 

tests/components/yeelight/__init__.py

@@ -8,7 +8,7 @@ from async_upnp_client.search import SsdpSearchListener
 from yeelight import BulbException, BulbType
 from yeelight.main import _MODEL_SPECS
 
-from homeassistant.components import yeelight as hass_yeelight
+from homeassistant.components import yeelight as hass_yeelight, zeroconf
 from homeassistant.components.yeelight import (
     CONF_MODE_MUSIC,
     CONF_NIGHTLIGHT_SWITCH_TYPE,
@@ -39,14 +39,14 @@ CAPABILITIES = {
 
 ID_DECIMAL = f"{int(ID, 16):08d}"
 
-ZEROCONF_DATA = {
-    "host": IP_ADDRESS,
-    "port": 54321,
-    "hostname": f"yeelink-light-strip1_miio{ID_DECIMAL}.local.",
-    "type": "_miio._udp.local.",
-    "name": f"yeelink-light-strip1_miio{ID_DECIMAL}._miio._udp.local.",
-    "properties": {"epoch": "1", "mac": "000000000000"},
-}
+ZEROCONF_DATA = zeroconf.ZeroconfServiceInfo(
+    host=IP_ADDRESS,
+    port=54321,
+    hostname=f"yeelink-light-strip1_miio{ID_DECIMAL}.local.",
+    type="_miio._udp.local.",
+    name=f"yeelink-light-strip1_miio{ID_DECIMAL}._miio._udp.local.",
+    properties={"epoch": "1", "mac": "000000000000"},
+)
 
 NAME = "name"
 SHORT_ID = hex(int("0x000000000015243f", 16))

tests/components/yeelight/test_config_flow.py

@@ -4,6 +4,7 @@ from unittest.mock import patch
 import pytest
 
 from homeassistant import config_entries
+from homeassistant.components import dhcp, zeroconf
 from homeassistant.components.yeelight import (
     CONF_DETECTED_MODEL,
     CONF_MODE_MUSIC,
@@ -455,7 +456,10 @@ async def test_discovered_by_homekit_and_dhcp(hass):
         result = await hass.config_entries.flow.async_init(
             DOMAIN,
             context={"source": config_entries.SOURCE_HOMEKIT},
-            data={"host": IP_ADDRESS, "properties": {"id": "aa:bb:cc:dd:ee:ff"}},
+            data=zeroconf.ZeroconfServiceInfo(
+                host=IP_ADDRESS,
+                properties={zeroconf.ATTR_PROPERTIES_ID: "aa:bb:cc:dd:ee:ff"},
+            ),
         )
         await hass.async_block_till_done()
     assert result["type"] == RESULT_TYPE_FORM
@@ -467,7 +471,7 @@ async def test_discovered_by_homekit_and_dhcp(hass):
         result2 = await hass.config_entries.flow.async_init(
             DOMAIN,
             context={"source": config_entries.SOURCE_DHCP},
-            data={"ip": IP_ADDRESS, "macaddress": "aa:bb:cc:dd:ee:ff"},
+            data=dhcp.DhcpServiceInfo(ip=IP_ADDRESS, macaddress="aa:bb:cc:dd:ee:ff"),
         )
         await hass.async_block_till_done()
     assert result2["type"] == RESULT_TYPE_ABORT
@@ -479,7 +483,7 @@ async def test_discovered_by_homekit_and_dhcp(hass):
         result3 = await hass.config_entries.flow.async_init(
             DOMAIN,
             context={"source": config_entries.SOURCE_DHCP},
-            data={"ip": IP_ADDRESS, "macaddress": "00:00:00:00:00:00"},
+            data=dhcp.DhcpServiceInfo(ip=IP_ADDRESS, macaddress="00:00:00:00:00:00"),
         )
         await hass.async_block_till_done()
     assert result3["type"] == RESULT_TYPE_ABORT
@@ -493,7 +497,7 @@ async def test_discovered_by_homekit_and_dhcp(hass):
         result3 = await hass.config_entries.flow.async_init(
             DOMAIN,
             context={"source": config_entries.SOURCE_DHCP},
-            data={"ip": "1.2.3.5", "macaddress": "00:00:00:00:00:01"},
+            data=dhcp.DhcpServiceInfo(ip="1.2.3.5", macaddress="00:00:00:00:00:01"),
         )
         await hass.async_block_till_done()
     assert result3["type"] == RESULT_TYPE_ABORT
@@ -505,11 +509,14 @@ async def test_discovered_by_homekit_and_dhcp(hass):
     [
         (
             config_entries.SOURCE_DHCP,
-            {"ip": IP_ADDRESS, "macaddress": "aa:bb:cc:dd:ee:ff"},
+            dhcp.DhcpServiceInfo(ip=IP_ADDRESS, macaddress="aa:bb:cc:dd:ee:ff"),
         ),
         (
             config_entries.SOURCE_HOMEKIT,
-            {"host": IP_ADDRESS, "properties": {"id": "aa:bb:cc:dd:ee:ff"}},
+            zeroconf.ZeroconfServiceInfo(
+                host=IP_ADDRESS,
+                properties={zeroconf.ATTR_PROPERTIES_ID: "aa:bb:cc:dd:ee:ff"},
+            ),
         ),
     ],
 )
@@ -563,11 +570,14 @@ async def test_discovered_by_dhcp_or_homekit(hass, source, data):
     [
         (
             config_entries.SOURCE_DHCP,
-            {"ip": IP_ADDRESS, "macaddress": "aa:bb:cc:dd:ee:ff"},
+            dhcp.DhcpServiceInfo(ip=IP_ADDRESS, macaddress="aa:bb:cc:dd:ee:ff"),
         ),
         (
             config_entries.SOURCE_HOMEKIT,
-            {"host": IP_ADDRESS, "properties": {"id": "aa:bb:cc:dd:ee:ff"}},
+            zeroconf.ZeroconfServiceInfo(
+                host=IP_ADDRESS,
+                properties={zeroconf.ATTR_PROPERTIES_ID: "aa:bb:cc:dd:ee:ff"},
+            ),
         ),
     ],
 )