From 73a4dba2ae6073cd3e37b37fa253050f30d896eb Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Wed, 1 Dec 2021 09:28:19 +0100
Subject: [PATCH] Use dataclass properties in yeelight discovery (#60640)

Co-authored-by: epenet <epenet@users.noreply.github.com>
---
 homeassistant/components/yeelight/config_flow.py | 16 ++++++++--------
 homeassistant/components/yeelight/manifest.json  |  3 ++-
 tests/components/yeelight/__init__.py            | 16 ++++++++++------
 tests/components/yeelight/test_config_flow.py    |  3 ++-
 4 files changed, 22 insertions(+), 16 deletions(-)

diff --git a/homeassistant/components/yeelight/config_flow.py b/homeassistant/components/yeelight/config_flow.py
index ba1ab4d24c7..0419824492a 100644
--- a/homeassistant/components/yeelight/config_flow.py
+++ b/homeassistant/components/yeelight/config_flow.py
@@ -8,7 +8,7 @@ from yeelight.aio import AsyncBulb
 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 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         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):
diff --git a/homeassistant/components/yeelight/manifest.json b/homeassistant/components/yeelight/manifest.json
index 13eb73cdd0c..60098514125 100644
--- a/homeassistant/components/yeelight/manifest.json
+++ b/homeassistant/components/yeelight/manifest.json
@@ -16,5 +16,6 @@
   "zeroconf": [{ "type": "_miio._udp.local.", "name": "yeelink-*" }],
   "homekit": {
     "models": ["YL*"]
-  }
+  },
+  "after_dependencies": ["ssdp"]
 }
diff --git a/tests/components/yeelight/__init__.py b/tests/components/yeelight/__init__.py
index a93cea7a94c..6f7ae807c9d 100644
--- a/tests/components/yeelight/__init__.py
+++ b/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 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",
diff --git a/tests/components/yeelight/test_config_flow.py b/tests/components/yeelight/test_config_flow.py
index b101fd3413d..aa5e7f98a45 100644
--- a/tests/components/yeelight/test_config_flow.py
+++ b/tests/components/yeelight/test_config_flow.py
@@ -55,7 +55,8 @@ DEFAULT_CONFIG = {
 SSDP_INFO = ssdp.SsdpServiceInfo(
     ssdp_usn="mock_usn",
     ssdp_st="mock_st",
-    upnp=CAPABILITIES,
+    upnp={},
+    ssdp_headers=CAPABILITIES,
 )