Use dataclass properties in yeelight discovery

homeassistant/components/yeelight/config_flow.py

@@ -8,7 +8,7 @@
 from yeelight.main import get_known_models
 
 from homeassistant import config_entries, exceptions
-from homeassistant.components import dhcp, zeroconf
+from homeassistant.components import dhcp, ssdp, zeroconf
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.const import CONF_DEVICE, CONF_HOST, CONF_ID, CONF_NAME
 from homeassistant.core import callback
@@ -60,28 +60,28 @@ async def async_step_homekit(
         self, discovery_info: zeroconf.ZeroconfServiceInfo
     ) -> FlowResult:
         """Handle discovery from homekit."""
-        self._discovered_ip = discovery_info[zeroconf.ATTR_HOST]
+        self._discovered_ip = discovery_info.host
         return await self._async_handle_discovery()
 
     async def async_step_dhcp(self, discovery_info: dhcp.DhcpServiceInfo) -> FlowResult:
         """Handle discovery from dhcp."""
-        self._discovered_ip = discovery_info[dhcp.IP_ADDRESS]
+        self._discovered_ip = discovery_info.ip
         return await self._async_handle_discovery()
 
     async def async_step_zeroconf(
         self, discovery_info: zeroconf.ZeroconfServiceInfo
     ) -> FlowResult:
         """Handle discovery from zeroconf."""
-        self._discovered_ip = discovery_info[zeroconf.ATTR_HOST]
+        self._discovered_ip = discovery_info.host
         await self.async_set_unique_id(
-            "{0:#0{1}x}".format(int(discovery_info[zeroconf.ATTR_NAME][-26:-18]), 18)
+            "{0:#0{1}x}".format(int(discovery_info.name[-26:-18]), 18)
         )
         return await self._async_handle_discovery_with_unique_id()
 
-    async def async_step_ssdp(self, discovery_info):
+    async def async_step_ssdp(self, discovery_info: ssdp.SsdpServiceInfo) -> FlowResult:
         """Handle discovery from ssdp."""
-        self._discovered_ip = urlparse(discovery_info["location"]).hostname
-        await self.async_set_unique_id(discovery_info["id"])
+        self._discovered_ip = urlparse(discovery_info.ssdp_headers["location"]).hostname
+        await self.async_set_unique_id(discovery_info.ssdp_headers["id"])
         return await self._async_handle_discovery_with_unique_id()
 
     async def _async_handle_discovery_with_unique_id(self):

homeassistant/components/yeelight/manifest.json

@@ -16,5 +16,6 @@
   "zeroconf": [{ "type": "_miio._udp.local.", "name": "yeelink-*" }],
   "homekit": {
     "models": ["YL*"]
-  }
+  },
+  "after_dependencies": ["ssdp"]
 }

tests/components/yeelight/__init__.py

@@ -8,7 +8,7 @@
 from yeelight import BulbException, BulbType
 from yeelight.main import _MODEL_SPECS
 
-from homeassistant.components import zeroconf
+from homeassistant.components import ssdp, zeroconf
 from homeassistant.components.yeelight import (
     CONF_MODE_MUSIC,
     CONF_NIGHTLIGHT_SWITCH_TYPE,
@@ -179,11 +179,15 @@ def _patch_discovery(no_device=False, capabilities=None):
     YeelightScanner._scanner = None  # Clear class scanner to reset hass
 
     def _generate_fake_ssdp_listener(*args, **kwargs):
-        return _patched_ssdp_listener(
-            None if no_device else capabilities or CAPABILITIES,
-            *args,
-            **kwargs,
-        )
+        info = None
+        if not no_device:
+            info = ssdp.SsdpServiceInfo(
+                ssdp_usn="mock_usn",
+                ssdp_st="mock_st",
+                upnp={},
+                ssdp_headers=capabilities or CAPABILITIES,
+            )
+        return _patched_ssdp_listener(info, *args, **kwargs)
 
     return patch(
         "homeassistant.components.yeelight.scanner.SsdpSearchListener",

tests/components/yeelight/test_config_flow.py

@@ -55,7 +55,8 @@
 SSDP_INFO = ssdp.SsdpServiceInfo(
     ssdp_usn="mock_usn",
     ssdp_st="mock_st",
-    upnp=CAPABILITIES,
+    upnp={},
+    ssdp_headers=CAPABILITIES,
 )