Make yeelight discovery async (#54711)

homeassistant/components/yeelight/__init__.py

@@ -2,13 +2,17 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 from datetime import timedelta
 import logging
+from urllib.parse import urlparse
 
+from async_upnp_client.search import SSDPListener
 import voluptuous as vol
-from yeelight import BulbException, discover_bulbs
+from yeelight import BulbException
 from yeelight.aio import KEY_CONNECTED, AsyncBulb
 
+from homeassistant import config_entries
 from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry, ConfigEntryNotReady
 from homeassistant.const import (
     CONF_DEVICES,
@@ -24,6 +28,7 @@ from homeassistant.helpers.dispatcher import (
     async_dispatcher_send,
 )
 from homeassistant.helpers.entity import DeviceInfo, Entity
+from homeassistant.helpers.event import async_call_later, async_track_time_interval
 from homeassistant.helpers.typing import ConfigType
 
 _LOGGER = logging.getLogger(__name__)
@@ -69,6 +74,12 @@ ACTIVE_COLOR_FLOWING = "1"
 NIGHTLIGHT_SWITCH_TYPE_LIGHT = "light"
 
 DISCOVERY_INTERVAL = timedelta(seconds=60)
+SSDP_TARGET = ("239.255.255.250", 1982)
+SSDP_ST = "wifi_bulb"
+DISCOVERY_ATTEMPTS = 3
+DISCOVERY_SEARCH_INTERVAL = timedelta(seconds=2)
+DISCOVERY_TIMEOUT = 2
+
 
 YEELIGHT_RGB_TRANSITION = "RGBTransition"
 YEELIGHT_HSV_TRANSACTION = "HSVTransition"
@@ -193,20 +204,12 @@ async def _async_initialize(
         hass.config_entries.async_setup_platforms(entry, PLATFORMS)
 
     if not device:
+        # get device and start listening for local pushes
         device = await _async_get_device(hass, host, entry)
+
+    await device.async_setup()
     entry_data[DATA_DEVICE] = device
 
-    # start listening for local pushes
-    await device.bulb.async_listen(device.async_update_callback)
-
-    # register stop callback to shutdown listening for local pushes
-    async def async_stop_listen_task(event):
-        """Stop listen thread."""
-        _LOGGER.debug("Shutting down Yeelight Listener")
-        await device.bulb.async_stop_listening()
-
-    hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, async_stop_listen_task)
-
     entry.async_on_unload(
         async_dispatcher_connect(
             hass, DEVICE_INITIALIZED.format(host), _async_load_platforms
@@ -251,7 +254,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     if entry.data.get(CONF_HOST):
         try:
             device = await _async_get_device(hass, entry.data[CONF_HOST], entry)
-        except OSError as ex:
+        except BulbException as ex:
             # If CONF_ID is not valid we cannot fallback to discovery
             # so we must retry by raising ConfigEntryNotReady
             if not entry.data.get(CONF_ID):
@@ -267,16 +270,16 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
                 raise ConfigEntryNotReady from ex
             return True
 
-    # discovery
+    async def _async_from_discovery(capabilities: dict[str, str]) -> None:
-    scanner = YeelightScanner.async_get(hass)
+        host = urlparse(capabilities["location"]).hostname
-
-    async def _async_from_discovery(host: str) -> None:
         try:
             await _async_initialize(hass, entry, host)
         except BulbException:
             _LOGGER.exception("Failed to connect to bulb at %s", host)
 
-    scanner.async_register_callback(entry.data[CONF_ID], _async_from_discovery)
+    # discovery
+    scanner = YeelightScanner.async_get(hass)
+    await scanner.async_register_callback(entry.data[CONF_ID], _async_from_discovery)
     return True
 
 
@@ -294,6 +297,7 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         scanner = YeelightScanner.async_get(hass)
         scanner.async_unregister_callback(entry.data[CONF_ID])
 
+    if DATA_DEVICE in entry_data:
         device = entry_data[DATA_DEVICE]
         _LOGGER.debug("Shutting down Yeelight Listener")
         await device.bulb.async_stop_listening()
@@ -307,9 +311,9 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 @callback
 def _async_unique_name(capabilities: dict) -> str:
     """Generate name from capabilities."""
-    model = capabilities["model"]
+    model = str(capabilities["model"]).replace("_", " ").title()
-    unique_id = capabilities["id"]
+    short_id = hex(int(capabilities["id"], 16))
-    return f"yeelight_{model}_{unique_id}"
+    return f"Yeelight {model} {short_id}"
 
 
 async def _async_update_listener(hass: HomeAssistant, entry: ConfigEntry):
@@ -333,88 +337,147 @@ class YeelightScanner:
     def __init__(self, hass: HomeAssistant) -> None:
         """Initialize class."""
         self._hass = hass
-        self._seen = {}
         self._callbacks = {}
-        self._scan_task = None
+        self._host_discovered_events = {}
+        self._unique_id_capabilities = {}
+        self._host_capabilities = {}
+        self._track_interval = None
+        self._listener = None
+        self._connected_event = None
 
-    async def _async_scan(self):
+    async def async_setup(self):
-        _LOGGER.debug("Yeelight scanning")
+        """Set up the scanner."""
-        # Run 3 times as packets can get lost
+        if self._connected_event:
-        for _ in range(3):
+            await self._connected_event.wait()
-            devices = await self._hass.async_add_executor_job(discover_bulbs)
+            return
-            for device in devices:
+        self._connected_event = asyncio.Event()
-                unique_id = device["capabilities"]["id"]
-                if unique_id in self._seen:
-                    continue
-                host = device["ip"]
-                self._seen[unique_id] = host
-                _LOGGER.debug("Yeelight discovered at %s", host)
-                if unique_id in self._callbacks:
-                    self._hass.async_create_task(self._callbacks[unique_id](host))
-                    self._callbacks.pop(unique_id)
-                    if len(self._callbacks) == 0:
-                        self._async_stop_scan()
 
-        await asyncio.sleep(DISCOVERY_INTERVAL.total_seconds())
+        async def _async_connected():
-        self._scan_task = self._hass.loop.create_task(self._async_scan())
+            self._listener.async_search()
+            self._connected_event.set()
+
+        self._listener = SSDPListener(
+            async_callback=self._async_process_entry,
+            service_type=SSDP_ST,
+            target=SSDP_TARGET,
+            async_connect_callback=_async_connected,
+        )
+        await self._listener.async_start()
+        await self._connected_event.wait()
+
+    async def async_discover(self):
+        """Discover bulbs."""
+        await self.async_setup()
+        for _ in range(DISCOVERY_ATTEMPTS):
+            self._listener.async_search()
+            await asyncio.sleep(DISCOVERY_SEARCH_INTERVAL.total_seconds())
+        return self._unique_id_capabilities.values()
 
     @callback
-    def _async_start_scan(self):
+    def async_scan(self, *_):
+        """Send discovery packets."""
+        _LOGGER.debug("Yeelight scanning")
+        self._listener.async_search()
+
+    async def async_get_capabilities(self, host):
+        """Get capabilities via SSDP."""
+        if host in self._host_capabilities:
+            return self._host_capabilities[host]
+
+        host_event = asyncio.Event()
+        self._host_discovered_events.setdefault(host, []).append(host_event)
+        await self.async_setup()
+
+        self._listener.async_search((host, SSDP_TARGET[1]))
+
+        with contextlib.suppress(asyncio.TimeoutError):
+            await asyncio.wait_for(host_event.wait(), timeout=DISCOVERY_TIMEOUT)
+
+        self._host_discovered_events[host].remove(host_event)
+        return self._host_capabilities.get(host)
+
+    def _async_discovered_by_ssdp(self, response):
+        @callback
+        def _async_start_flow(*_):
+            asyncio.create_task(
+                self._hass.config_entries.flow.async_init(
+                    DOMAIN,
+                    context={"source": config_entries.SOURCE_SSDP},
+                    data=response,
+                )
+            )
+
+        # Delay starting the flow in case the discovery is the result
+        # of another discovery
+        async_call_later(self._hass, 1, _async_start_flow)
+
+    async def _async_process_entry(self, response):
+        """Process a discovery."""
+        _LOGGER.debug("Discovered via SSDP: %s", response)
+        unique_id = response["id"]
+        host = urlparse(response["location"]).hostname
+        if unique_id not in self._unique_id_capabilities:
+            _LOGGER.debug("Yeelight discovered with %s", response)
+            self._async_discovered_by_ssdp(response)
+        self._host_capabilities[host] = response
+        self._unique_id_capabilities[unique_id] = response
+        for event in self._host_discovered_events.get(host, []):
+            event.set()
+        if unique_id in self._callbacks:
+            self._hass.async_create_task(self._callbacks[unique_id](response))
+            self._callbacks.pop(unique_id)
+        if not self._callbacks:
+            self._async_stop_scan()
+
+    async def _async_start_scan(self):
         """Start scanning for Yeelight devices."""
         _LOGGER.debug("Start scanning")
-        # Use loop directly to avoid home assistant track this task
+        await self.async_setup()
-        self._scan_task = self._hass.loop.create_task(self._async_scan())
+        if not self._track_interval:
+            self._track_interval = async_track_time_interval(
+                self._hass, self.async_scan, DISCOVERY_INTERVAL
+            )
+            self.async_scan()
 
     @callback
     def _async_stop_scan(self):
         """Stop scanning."""
-        _LOGGER.debug("Stop scanning")
+        if self._track_interval is None:
-        if self._scan_task is not None:
+            return
-            self._scan_task.cancel()
+        _LOGGER.debug("Stop scanning interval")
-            self._scan_task = None
+        self._track_interval()
+        self._track_interval = None
 
-    @callback
+    async def async_register_callback(self, unique_id, callback_func):
-    def async_register_callback(self, unique_id, callback_func):
         """Register callback function."""
-        host = self._seen.get(unique_id)
+        if capabilities := self._unique_id_capabilities.get(unique_id):
-        if host is not None:
+            self._hass.async_create_task(callback_func(capabilities))
-            self._hass.async_create_task(callback_func(host))
+            return
-        else:
         self._callbacks[unique_id] = callback_func
-            if len(self._callbacks) == 1:
+        await self._async_start_scan()
-                self._async_start_scan()
 
     @callback
     def async_unregister_callback(self, unique_id):
         """Unregister callback function."""
-        if unique_id not in self._callbacks:
+        self._callbacks.pop(unique_id, None)
-            return
+        if not self._callbacks:
-        self._callbacks.pop(unique_id)
-        if len(self._callbacks) == 0:
             self._async_stop_scan()
 
 
 class YeelightDevice:
     """Represents single Yeelight device."""
 
-    def __init__(self, hass, host, config, bulb, capabilities):
+    def __init__(self, hass, host, config, bulb):
         """Initialize device."""
         self._hass = hass
         self._config = config
         self._host = host
         self._bulb_device = bulb
-        self._capabilities = capabilities or {}
+        self._capabilities = {}
         self._device_type = None
         self._available = False
         self._initialized = False
-
+        self._name = None
-        self._name = host  # Default name is host
-        if capabilities:
-            # Generate name from model and id when capabilities is available
-            self._name = _async_unique_name(capabilities)
-        if config.get(CONF_NAME):
-            # Override default name when name is set in config
-            self._name = config[CONF_NAME]
 
     @property
     def bulb(self):
@@ -444,7 +507,7 @@ class YeelightDevice:
     @property
     def model(self):
         """Return configured/autodetected device model."""
-        return self._bulb_device.model
+        return self._bulb_device.model or self._capabilities.get("model")
 
     @property
     def fw_version(self):
@@ -530,7 +593,8 @@ class YeelightDevice:
             await self.bulb.async_get_properties(UPDATE_REQUEST_PROPERTIES)
             self._available = True
             if not self._initialized:
-                await self._async_initialize_device()
+                self._initialized = True
+                async_dispatcher_send(self._hass, DEVICE_INITIALIZED.format(self._host))
         except BulbException as ex:
             if self._available:  # just inform once
                 _LOGGER.error(
@@ -540,28 +604,18 @@ class YeelightDevice:
 
         return self._available
 
-    async def _async_get_capabilities(self):
+    async def async_setup(self):
-        """Request device capabilities."""
+        """Fetch capabilities and setup name if available."""
-        try:
+        scanner = YeelightScanner.async_get(self._hass)
-            await self._hass.async_add_executor_job(self.bulb.get_capabilities)
+        self._capabilities = await scanner.async_get_capabilities(self._host) or {}
-            _LOGGER.debug(
+        if name := self._config.get(CONF_NAME):
-                "Device %s, %s capabilities: %s",
+            # Override default name when name is set in config
-                self._host,
+            self._name = name
-                self.name,
+        elif self._capabilities:
-                self.bulb.capabilities,
+            # Generate name from model and id when capabilities is available
-            )
+            self._name = _async_unique_name(self._capabilities)
-        except BulbException as ex:
+        else:
-            _LOGGER.error(
+            self._name = self._host  # Default name is host
-                "Unable to get device capabilities %s, %s: %s",
-                self._host,
-                self.name,
-                ex,
-            )
-
-    async def _async_initialize_device(self):
-        await self._async_get_capabilities()
-        self._initialized = True
-        async_dispatcher_send(self._hass, DEVICE_INITIALIZED.format(self._host))
 
     async def async_update(self):
         """Update device properties and send data updated signal."""
@@ -628,6 +682,19 @@ async def _async_get_device(
 
     # Set up device
     bulb = AsyncBulb(host, model=model or None)
-    capabilities = await hass.async_add_executor_job(bulb.get_capabilities)
 
-    return YeelightDevice(hass, host, entry.options, bulb, capabilities)
+    device = YeelightDevice(hass, host, entry.options, bulb)
+    # start listening for local pushes
+    await device.bulb.async_listen(device.async_update_callback)
+
+    # register stop callback to shutdown listening for local pushes
+    async def async_stop_listen_task(event):
+        """Stop listen thread."""
+        _LOGGER.debug("Shutting down Yeelight Listener")
+        await device.bulb.async_stop_listening()
+
+    entry.async_on_unload(
+        hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, async_stop_listen_task)
+    )
+
+    return device

homeassistant/components/yeelight/config_flow.py

@@ -1,8 +1,10 @@
 """Config flow for Yeelight integration."""
 import logging
+from urllib.parse import urlparse
 
 import voluptuous as vol
 import yeelight
+from yeelight.aio import AsyncBulb
 
 from homeassistant import config_entries, exceptions
 from homeassistant.components.dhcp import IP_ADDRESS
@@ -19,6 +21,7 @@ from . import (
     CONF_TRANSITION,
     DOMAIN,
     NIGHTLIGHT_SWITCH_TYPE_LIGHT,
+    YeelightScanner,
     _async_unique_name,
 )
 
@@ -54,6 +57,15 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         self._discovered_ip = discovery_info[IP_ADDRESS]
         return await self._async_handle_discovery()
 
+    async def async_step_ssdp(self, discovery_info):
+        """Handle discovery from ssdp."""
+        self._discovered_ip = urlparse(discovery_info["location"]).hostname
+        await self.async_set_unique_id(discovery_info["id"])
+        self._abort_if_unique_id_configured(
+            updates={CONF_HOST: self._discovered_ip}, reload_on_update=False
+        )
+        return await self._async_handle_discovery()
+
     async def _async_handle_discovery(self):
         """Handle any discovery."""
         self.context[CONF_HOST] = self._discovered_ip
@@ -62,7 +74,9 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
                 return self.async_abort(reason="already_in_progress")
 
         try:
-            self._discovered_model = await self._async_try_connect(self._discovered_ip)
+            self._discovered_model = await self._async_try_connect(
+                self._discovered_ip, raise_on_progress=True
+            )
         except CannotConnect:
             return self.async_abort(reason="cannot_connect")
 
@@ -96,7 +110,9 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
             if not user_input.get(CONF_HOST):
                 return await self.async_step_pick_device()
             try:
-                model = await self._async_try_connect(user_input[CONF_HOST])
+                model = await self._async_try_connect(
+                    user_input[CONF_HOST], raise_on_progress=False
+                )
             except CannotConnect:
                 errors["base"] = "cannot_connect"
             else:
@@ -119,10 +135,12 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         if user_input is not None:
             unique_id = user_input[CONF_DEVICE]
             capabilities = self._discovered_devices[unique_id]
-            await self.async_set_unique_id(unique_id)
+            await self.async_set_unique_id(unique_id, raise_on_progress=False)
             self._abort_if_unique_id_configured()
+            host = urlparse(capabilities["location"]).hostname
             return self.async_create_entry(
-                title=_async_unique_name(capabilities), data={CONF_ID: unique_id}
+                title=_async_unique_name(capabilities),
+                data={CONF_ID: unique_id, CONF_HOST: host},
             )
 
         configured_devices = {
@@ -131,16 +149,15 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
             if entry.data[CONF_ID]
         }
         devices_name = {}
+        scanner = YeelightScanner.async_get(self.hass)
+        devices = await scanner.async_discover()
         # Run 3 times as packets can get lost
-        for _ in range(3):
+        for capabilities in devices:
-            devices = await self.hass.async_add_executor_job(yeelight.discover_bulbs)
-            for device in devices:
-                capabilities = device["capabilities"]
             unique_id = capabilities["id"]
             if unique_id in configured_devices:
                 continue  # ignore configured devices
             model = capabilities["model"]
-                host = device["ip"]
+            host = urlparse(capabilities["location"]).hostname
             name = f"{host} {model} {unique_id}"
             self._discovered_devices[unique_id] = capabilities
             devices_name[unique_id] = name
@@ -157,7 +174,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         """Handle import step."""
         host = user_input[CONF_HOST]
         try:
-            await self._async_try_connect(host)
+            await self._async_try_connect(host, raise_on_progress=False)
         except CannotConnect:
             _LOGGER.error("Failed to import %s: cannot connect", host)
             return self.async_abort(reason="cannot_connect")
@@ -169,27 +186,26 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         self._abort_if_unique_id_configured()
         return self.async_create_entry(title=user_input[CONF_NAME], data=user_input)
 
-    async def _async_try_connect(self, host):
+    async def _async_try_connect(self, host, raise_on_progress=True):
         """Set up with options."""
         self._async_abort_entries_match({CONF_HOST: host})
 
-        bulb = yeelight.Bulb(host)
+        scanner = YeelightScanner.async_get(self.hass)
-        try:
+        capabilities = await scanner.async_get_capabilities(host)
-            capabilities = await self.hass.async_add_executor_job(bulb.get_capabilities)
         if capabilities is None:  # timeout
             _LOGGER.debug("Failed to get capabilities from %s: timeout", host)
         else:
             _LOGGER.debug("Get capabilities: %s", capabilities)
-                await self.async_set_unique_id(capabilities["id"])
+            await self.async_set_unique_id(
+                capabilities["id"], raise_on_progress=raise_on_progress
+            )
             return capabilities["model"]
-        except OSError as err:
-            _LOGGER.debug("Failed to get capabilities from %s: %s", host, err)
-            # Ignore the error since get_capabilities uses UDP discovery packet
-            # which does not work in all network environments
-
         # Fallback to get properties
+        bulb = AsyncBulb(host)
         try:
-            await self.hass.async_add_executor_job(bulb.get_properties)
+            await bulb.async_listen(lambda _: True)
+            await bulb.async_get_properties()
+            await bulb.async_stop_listening()
         except yeelight.BulbException as err:
             _LOGGER.error("Failed to get properties from %s: %s", host, err)
             raise CannotConnect from err

homeassistant/components/yeelight/light.py

@@ -905,7 +905,7 @@ class YeelightNightLightMode(YeelightGenericLight):
     @property
     def name(self) -> str:
         """Return the name of the device if any."""
-        return f"{self.device.name} nightlight"
+        return f"{self.device.name} Nightlight"
 
     @property
     def icon(self):
@@ -997,7 +997,7 @@ class YeelightAmbientLight(YeelightColorLightWithoutNightlightSwitch):
     @property
     def name(self) -> str:
         """Return the name of the device if any."""
-        return f"{self.device.name} ambilight"
+        return f"{self.device.name} Ambilight"
 
     @property
     def _brightness_property(self):

homeassistant/components/yeelight/manifest.json

@@ -2,7 +2,7 @@
   "domain": "yeelight",
   "name": "Yeelight",
   "documentation": "https://www.home-assistant.io/integrations/yeelight",
-  "requirements": ["yeelight==0.7.2"],
+  "requirements": ["yeelight==0.7.2", "async-upnp-client==0.20.0"],
   "codeowners": ["@rytilahti", "@zewelor", "@shenxn", "@starkillerOG"],
   "config_flow": true,
   "iot_class": "local_push",

requirements_all.txt

@@ -314,6 +314,7 @@ asterisk_mbox==0.5.0
 # homeassistant.components.dlna_dmr
 # homeassistant.components.ssdp
 # homeassistant.components.upnp
+# homeassistant.components.yeelight
 async-upnp-client==0.20.0
 
 # homeassistant.components.supla

requirements_test_all.txt

@@ -205,6 +205,7 @@ arcam-fmj==0.7.0
 # homeassistant.components.dlna_dmr
 # homeassistant.components.ssdp
 # homeassistant.components.upnp
+# homeassistant.components.yeelight
 async-upnp-client==0.20.0
 
 # homeassistant.components.aurora

tests/components/yeelight/__init__.py

@@ -1,9 +1,13 @@
 """Tests for the Yeelight integration."""
+import asyncio
+from datetime import timedelta
 from unittest.mock import AsyncMock, MagicMock, patch
 
+from async_upnp_client.search import SSDPListener
 from yeelight import BulbException, BulbType
 from yeelight.main import _MODEL_SPECS
 
+from homeassistant.components import yeelight as hass_yeelight
 from homeassistant.components.yeelight import (
     CONF_MODE_MUSIC,
     CONF_NIGHTLIGHT_SWITCH_TYPE,
@@ -13,6 +17,7 @@ from homeassistant.components.yeelight import (
     YeelightScanner,
 )
 from homeassistant.const import CONF_DEVICES, CONF_ID, CONF_NAME
+from homeassistant.core import callback
 
 IP_ADDRESS = "192.168.1.239"
 MODEL = "color"
@@ -23,13 +28,16 @@ CAPABILITIES = {
     "id": ID,
     "model": MODEL,
     "fw_ver": FW_VER,
+    "location": f"yeelight://{IP_ADDRESS}",
     "support": "get_prop set_default set_power toggle set_bright start_cf stop_cf"
     " set_scene cron_add cron_get cron_del set_ct_abx set_rgb",
     "name": "",
 }
 
 NAME = "name"
-UNIQUE_NAME = f"yeelight_{MODEL}_{ID}"
+SHORT_ID = hex(int("0x000000000015243f", 16))
+UNIQUE_NAME = f"yeelight_{MODEL}_{SHORT_ID}"
+UNIQUE_FRIENDLY_NAME = f"Yeelight {MODEL.title()} {SHORT_ID}"
 
 MODULE = "homeassistant.components.yeelight"
 MODULE_CONFIG_FLOW = f"{MODULE}.config_flow"
@@ -81,8 +89,8 @@ CONFIG_ENTRY_DATA = {CONF_ID: ID}
 
 def _mocked_bulb(cannot_connect=False):
     bulb = MagicMock()
-    type(bulb).get_capabilities = MagicMock(
+    type(bulb).async_listen = AsyncMock(
-        return_value=None if cannot_connect else CAPABILITIES
+        side_effect=BulbException if cannot_connect else None
     )
     type(bulb).async_get_properties = AsyncMock(
         side_effect=BulbException if cannot_connect else None
@@ -98,7 +106,6 @@ def _mocked_bulb(cannot_connect=False):
     bulb.last_properties = PROPERTIES.copy()
     bulb.music_mode = False
     bulb.async_get_properties = AsyncMock()
-    bulb.async_listen = AsyncMock()
     bulb.async_stop_listening = AsyncMock()
     bulb.async_update = AsyncMock()
     bulb.async_turn_on = AsyncMock()
@@ -116,12 +123,43 @@ def _mocked_bulb(cannot_connect=False):
     return bulb
 
 
-def _patch_discovery(prefix, no_device=False):
+def _patched_ssdp_listener(info, *args, **kwargs):
+    listener = SSDPListener(*args, **kwargs)
+
+    async def _async_callback(*_):
+        await listener.async_connect_callback()
+
+    @callback
+    def _async_search(*_):
+        if info:
+            asyncio.create_task(listener.async_callback(info))
+
+    listener.async_start = _async_callback
+    listener.async_search = _async_search
+    return listener
+
+
+def _patch_discovery(no_device=False):
     YeelightScanner._scanner = None  # Clear class scanner to reset hass
 
-    def _mocked_discovery(timeout=2, interface=False):
+    def _generate_fake_ssdp_listener(*args, **kwargs):
-        if no_device:
+        return _patched_ssdp_listener(
-            return []
+            None if no_device else CAPABILITIES,
-        return [{"ip": IP_ADDRESS, "port": 55443, "capabilities": CAPABILITIES}]
+            *args,
+            **kwargs,
+        )
 
-    return patch(f"{prefix}.discover_bulbs", side_effect=_mocked_discovery)
+    return patch(
+        "homeassistant.components.yeelight.SSDPListener",
+        new=_generate_fake_ssdp_listener,
+    )
+
+
+def _patch_discovery_interval():
+    return patch.object(
+        hass_yeelight, "DISCOVERY_SEARCH_INTERVAL", timedelta(seconds=0)
+    )
+
+
+def _patch_discovery_timeout():
+    return patch.object(hass_yeelight, "DISCOVERY_TIMEOUT", 0.0001)

tests/components/yeelight/test_binary_sensor.py

@@ -6,7 +6,14 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_component
 from homeassistant.setup import async_setup_component
 
-from . import MODULE, NAME, PROPERTIES, YAML_CONFIGURATION, _mocked_bulb
+from . import (
+    MODULE,
+    NAME,
+    PROPERTIES,
+    YAML_CONFIGURATION,
+    _mocked_bulb,
+    _patch_discovery,
+)
 
 ENTITY_BINARY_SENSOR = f"binary_sensor.{NAME}_nightlight"
 
@@ -14,9 +21,7 @@ ENTITY_BINARY_SENSOR = f"binary_sensor.{NAME}_nightlight"
 async def test_nightlight(hass: HomeAssistant):
     """Test nightlight sensor."""
     mocked_bulb = _mocked_bulb()
-    with patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb), patch(
+    with _patch_discovery(), patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb):
-        f"{MODULE}.config_flow.yeelight.Bulb", return_value=mocked_bulb
-    ):
         await async_setup_component(hass, DOMAIN, YAML_CONFIGURATION)
         await hass.async_block_till_done()
 

tests/components/yeelight/test_config_flow.py

@@ -1,5 +1,5 @@
 """Test the Yeelight config flow."""
-from unittest.mock import MagicMock, patch
+from unittest.mock import patch
 
 import pytest
 
@@ -25,14 +25,17 @@ from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import RESULT_TYPE_ABORT, RESULT_TYPE_FORM
 
 from . import (
+    CAPABILITIES,
     ID,
     IP_ADDRESS,
     MODULE,
     MODULE_CONFIG_FLOW,
     NAME,
-    UNIQUE_NAME,
+    UNIQUE_FRIENDLY_NAME,
     _mocked_bulb,
     _patch_discovery,
+    _patch_discovery_interval,
+    _patch_discovery_timeout,
 )
 
 from tests.common import MockConfigEntry
@@ -55,21 +58,23 @@ async def test_discovery(hass: HomeAssistant):
     assert result["step_id"] == "user"
     assert not result["errors"]
 
-    with _patch_discovery(f"{MODULE_CONFIG_FLOW}.yeelight"):
+    with _patch_discovery(), _patch_discovery_interval():
         result2 = await hass.config_entries.flow.async_configure(result["flow_id"], {})
     assert result2["type"] == "form"
     assert result2["step_id"] == "pick_device"
     assert not result2["errors"]
 
-    with patch(f"{MODULE}.async_setup", return_value=True) as mock_setup, patch(
+    with _patch_discovery(), _patch_discovery_interval(), patch(
+        f"{MODULE}.async_setup", return_value=True
+    ) as mock_setup, patch(
         f"{MODULE}.async_setup_entry", return_value=True
     ) as mock_setup_entry:
         result3 = await hass.config_entries.flow.async_configure(
             result["flow_id"], {CONF_DEVICE: ID}
         )
     assert result3["type"] == "create_entry"
-    assert result3["title"] == UNIQUE_NAME
+    assert result3["title"] == UNIQUE_FRIENDLY_NAME
-    assert result3["data"] == {CONF_ID: ID}
+    assert result3["data"] == {CONF_ID: ID, CONF_HOST: IP_ADDRESS}
     await hass.async_block_till_done()
     mock_setup.assert_called_once()
     mock_setup_entry.assert_called_once()
@@ -82,7 +87,7 @@ async def test_discovery(hass: HomeAssistant):
     assert result["step_id"] == "user"
     assert not result["errors"]
 
-    with _patch_discovery(f"{MODULE_CONFIG_FLOW}.yeelight"):
+    with _patch_discovery(), _patch_discovery_interval():
         result2 = await hass.config_entries.flow.async_configure(result["flow_id"], {})
     assert result2["type"] == "abort"
     assert result2["reason"] == "no_devices_found"
@@ -94,7 +99,9 @@ async def test_discovery_no_device(hass: HomeAssistant):
         DOMAIN, context={"source": config_entries.SOURCE_USER}
     )
 
-    with _patch_discovery(f"{MODULE_CONFIG_FLOW}.yeelight", no_device=True):
+    with _patch_discovery(
+        no_device=True
+    ), _patch_discovery_timeout(), _patch_discovery_interval():
         result2 = await hass.config_entries.flow.async_configure(result["flow_id"], {})
 
     assert result2["type"] == "abort"
@@ -114,26 +121,27 @@ async def test_import(hass: HomeAssistant):
 
     # Cannot connect
     mocked_bulb = _mocked_bulb(cannot_connect=True)
-    with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb):
+    with _patch_discovery(
+        no_device=True
+    ), _patch_discovery_timeout(), _patch_discovery_interval(), patch(
+        f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb
+    ):
         result = await hass.config_entries.flow.async_init(
             DOMAIN, context={"source": config_entries.SOURCE_IMPORT}, data=config
         )
-    type(mocked_bulb).get_capabilities.assert_called_once()
-    type(mocked_bulb).get_properties.assert_called_once()
     assert result["type"] == "abort"
     assert result["reason"] == "cannot_connect"
 
     # Success
     mocked_bulb = _mocked_bulb()
-    with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb), patch(
+    with _patch_discovery(), patch(
-        f"{MODULE}.async_setup", return_value=True
+        f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb
-    ) as mock_setup, patch(
+    ), patch(f"{MODULE}.async_setup", return_value=True) as mock_setup, patch(
         f"{MODULE}.async_setup_entry", return_value=True
     ) as mock_setup_entry:
         result = await hass.config_entries.flow.async_init(
             DOMAIN, context={"source": config_entries.SOURCE_IMPORT}, data=config
         )
-    type(mocked_bulb).get_capabilities.assert_called_once()
     assert result["type"] == "create_entry"
     assert result["title"] == DEFAULT_NAME
     assert result["data"] == {
@@ -150,7 +158,9 @@ async def test_import(hass: HomeAssistant):
 
     # Duplicate
     mocked_bulb = _mocked_bulb()
-    with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb):
+    with _patch_discovery(), patch(
+        f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb
+    ):
         result = await hass.config_entries.flow.async_init(
             DOMAIN, context={"source": config_entries.SOURCE_IMPORT}, data=config
         )
@@ -169,7 +179,11 @@ async def test_manual(hass: HomeAssistant):
 
     # Cannot connect (timeout)
     mocked_bulb = _mocked_bulb(cannot_connect=True)
-    with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb):
+    with _patch_discovery(
+        no_device=True
+    ), _patch_discovery_timeout(), _patch_discovery_interval(), patch(
+        f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb
+    ):
         result2 = await hass.config_entries.flow.async_configure(
             result["flow_id"], {CONF_HOST: IP_ADDRESS}
         )
@@ -178,8 +192,11 @@ async def test_manual(hass: HomeAssistant):
     assert result2["errors"] == {"base": "cannot_connect"}
 
     # Cannot connect (error)
-    type(mocked_bulb).get_capabilities = MagicMock(side_effect=OSError)
+    with _patch_discovery(
-    with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb):
+        no_device=True
+    ), _patch_discovery_timeout(), _patch_discovery_interval(), patch(
+        f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb
+    ):
         result3 = await hass.config_entries.flow.async_configure(
             result["flow_id"], {CONF_HOST: IP_ADDRESS}
         )
@@ -187,9 +204,11 @@ async def test_manual(hass: HomeAssistant):
 
     # Success
     mocked_bulb = _mocked_bulb()
-    with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb), patch(
+    with _patch_discovery(), _patch_discovery_timeout(), patch(
-        f"{MODULE}.async_setup", return_value=True
+        f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb
-    ), patch(f"{MODULE}.async_setup_entry", return_value=True):
+    ), patch(f"{MODULE}.async_setup", return_value=True), patch(
+        f"{MODULE}.async_setup_entry", return_value=True
+    ):
         result4 = await hass.config_entries.flow.async_configure(
             result["flow_id"], {CONF_HOST: IP_ADDRESS}
         )
@@ -203,7 +222,11 @@ async def test_manual(hass: HomeAssistant):
         DOMAIN, context={"source": config_entries.SOURCE_USER}
     )
     mocked_bulb = _mocked_bulb()
-    with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb):
+    with _patch_discovery(
+        no_device=True
+    ), _patch_discovery_timeout(), _patch_discovery_interval(), patch(
+        f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb
+    ):
         result2 = await hass.config_entries.flow.async_configure(
             result["flow_id"], {CONF_HOST: IP_ADDRESS}
         )
@@ -219,7 +242,7 @@ async def test_options(hass: HomeAssistant):
     config_entry.add_to_hass(hass)
 
     mocked_bulb = _mocked_bulb()
-    with patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb):
+    with _patch_discovery(), patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb):
         assert await hass.config_entries.async_setup(config_entry.entry_id)
         await hass.async_block_till_done()
 
@@ -241,7 +264,7 @@ async def test_options(hass: HomeAssistant):
     config[CONF_NIGHTLIGHT_SWITCH] = True
     user_input = {**config}
     user_input.pop(CONF_NAME)
-    with patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb):
+    with _patch_discovery(), patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb):
         result2 = await hass.config_entries.options.async_configure(
             result["flow_id"], user_input
         )
@@ -262,15 +285,18 @@ async def test_manual_no_capabilities(hass: HomeAssistant):
     assert not result["errors"]
 
     mocked_bulb = _mocked_bulb()
-    type(mocked_bulb).get_capabilities = MagicMock(return_value=None)
+    with _patch_discovery(
-    with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb), patch(
+        no_device=True
+    ), _patch_discovery_timeout(), _patch_discovery_interval(), patch(
+        f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb
+    ), patch(
         f"{MODULE}.async_setup", return_value=True
-    ), patch(f"{MODULE}.async_setup_entry", return_value=True):
+    ), patch(
+        f"{MODULE}.async_setup_entry", return_value=True
+    ):
         result = await hass.config_entries.flow.async_configure(
             result["flow_id"], {CONF_HOST: IP_ADDRESS}
         )
-    type(mocked_bulb).get_capabilities.assert_called_once()
-    type(mocked_bulb).get_properties.assert_called_once()
     assert result["type"] == "create_entry"
     assert result["data"] == {CONF_HOST: IP_ADDRESS}
 
@@ -280,39 +306,53 @@ async def test_discovered_by_homekit_and_dhcp(hass):
     await setup.async_setup_component(hass, "persistent_notification", {})
 
     mocked_bulb = _mocked_bulb()
-    with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb):
+    with _patch_discovery(), _patch_discovery_interval(), patch(
+        f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb
+    ):
         result = await hass.config_entries.flow.async_init(
             DOMAIN,
             context={"source": config_entries.SOURCE_HOMEKIT},
-            data={"host": "1.2.3.4", "properties": {"id": "aa:bb:cc:dd:ee:ff"}},
+            data={"host": IP_ADDRESS, "properties": {"id": "aa:bb:cc:dd:ee:ff"}},
         )
+        await hass.async_block_till_done()
     assert result["type"] == RESULT_TYPE_FORM
     assert result["errors"] is None
 
-    with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb):
+    with _patch_discovery(), _patch_discovery_interval(), patch(
+        f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb
+    ):
         result2 = await hass.config_entries.flow.async_init(
             DOMAIN,
             context={"source": config_entries.SOURCE_DHCP},
-            data={"ip": "1.2.3.4", "macaddress": "aa:bb:cc:dd:ee:ff"},
+            data={"ip": IP_ADDRESS, "macaddress": "aa:bb:cc:dd:ee:ff"},
         )
+        await hass.async_block_till_done()
     assert result2["type"] == RESULT_TYPE_ABORT
     assert result2["reason"] == "already_in_progress"
 
-    with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb):
+    with _patch_discovery(), _patch_discovery_interval(), patch(
+        f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb
+    ):
         result3 = await hass.config_entries.flow.async_init(
             DOMAIN,
             context={"source": config_entries.SOURCE_DHCP},
-            data={"ip": "1.2.3.4", "macaddress": "00:00:00:00:00:00"},
+            data={"ip": IP_ADDRESS, "macaddress": "00:00:00:00:00:00"},
         )
+        await hass.async_block_till_done()
     assert result3["type"] == RESULT_TYPE_ABORT
     assert result3["reason"] == "already_in_progress"
 
-    with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", side_effect=CannotConnect):
+    with _patch_discovery(
+        no_device=True
+    ), _patch_discovery_timeout(), _patch_discovery_interval(), patch(
+        f"{MODULE_CONFIG_FLOW}.AsyncBulb", side_effect=CannotConnect
+    ):
         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"},
         )
+        await hass.async_block_till_done()
     assert result3["type"] == RESULT_TYPE_ABORT
     assert result3["reason"] == "cannot_connect"
 
@@ -335,17 +375,25 @@ async def test_discovered_by_dhcp_or_homekit(hass, source, data):
     await setup.async_setup_component(hass, "persistent_notification", {})
 
     mocked_bulb = _mocked_bulb()
-    with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb):
+    with _patch_discovery(), _patch_discovery_interval(), patch(
+        f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb
+    ):
         result = await hass.config_entries.flow.async_init(
             DOMAIN, context={"source": source}, data=data
         )
+        await hass.async_block_till_done()
+
     assert result["type"] == RESULT_TYPE_FORM
     assert result["errors"] is None
 
-    with patch(f"{MODULE}.async_setup", return_value=True) as mock_async_setup, patch(
+    with _patch_discovery(), _patch_discovery_interval(), patch(
+        f"{MODULE}.async_setup", return_value=True
+    ) as mock_async_setup, patch(
         f"{MODULE}.async_setup_entry", return_value=True
     ) as mock_async_setup_entry:
         result2 = await hass.config_entries.flow.async_configure(result["flow_id"], {})
+        await hass.async_block_till_done()
+
     assert result2["type"] == "create_entry"
     assert result2["data"] == {CONF_HOST: IP_ADDRESS, CONF_ID: "0x000000000015243f"}
     assert mock_async_setup.called
@@ -370,10 +418,55 @@ async def test_discovered_by_dhcp_or_homekit_failed_to_get_id(hass, source, data
     await setup.async_setup_component(hass, "persistent_notification", {})
 
     mocked_bulb = _mocked_bulb()
-    type(mocked_bulb).get_capabilities = MagicMock(return_value=None)
+    with _patch_discovery(
-    with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb):
+        no_device=True
+    ), _patch_discovery_timeout(), _patch_discovery_interval(), patch(
+        f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb
+    ):
         result = await hass.config_entries.flow.async_init(
             DOMAIN, context={"source": source}, data=data
         )
     assert result["type"] == RESULT_TYPE_ABORT
     assert result["reason"] == "cannot_connect"
+
+
+async def test_discovered_ssdp(hass):
+    """Test we can setup when discovered from ssdp."""
+    await setup.async_setup_component(hass, "persistent_notification", {})
+
+    mocked_bulb = _mocked_bulb()
+    with _patch_discovery(), _patch_discovery_interval(), patch(
+        f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb
+    ):
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN, context={"source": config_entries.SOURCE_SSDP}, data=CAPABILITIES
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == RESULT_TYPE_FORM
+    assert result["errors"] is None
+
+    with _patch_discovery(), _patch_discovery_interval(), patch(
+        f"{MODULE}.async_setup", return_value=True
+    ) as mock_async_setup, patch(
+        f"{MODULE}.async_setup_entry", return_value=True
+    ) as mock_async_setup_entry:
+        result2 = await hass.config_entries.flow.async_configure(result["flow_id"], {})
+        await hass.async_block_till_done()
+
+    assert result2["type"] == "create_entry"
+    assert result2["data"] == {CONF_HOST: IP_ADDRESS, CONF_ID: "0x000000000015243f"}
+    assert mock_async_setup.called
+    assert mock_async_setup_entry.called
+
+    mocked_bulb = _mocked_bulb()
+    with _patch_discovery(), _patch_discovery_interval(), patch(
+        f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb
+    ):
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN, context={"source": config_entries.SOURCE_SSDP}, data=CAPABILITIES
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == RESULT_TYPE_ABORT
+    assert result["reason"] == "already_configured"

tests/components/yeelight/test_init.py

@@ -1,5 +1,6 @@
 """Test Yeelight."""
-from unittest.mock import AsyncMock, MagicMock, patch
+from datetime import timedelta
+from unittest.mock import AsyncMock, patch
 
 from yeelight import BulbException, BulbType
 
@@ -22,9 +23,9 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
 from homeassistant.setup import async_setup_component
+from homeassistant.util import dt as dt_util
 
 from . import (
-    CAPABILITIES,
     CONFIG_ENTRY_DATA,
     ENTITY_AMBILIGHT,
     ENTITY_BINARY_SENSOR,
@@ -34,12 +35,14 @@ from . import (
     ID,
     IP_ADDRESS,
     MODULE,
-    MODULE_CONFIG_FLOW,
+    SHORT_ID,
     _mocked_bulb,
     _patch_discovery,
+    _patch_discovery_interval,
+    _patch_discovery_timeout,
 )
 
-from tests.common import MockConfigEntry
+from tests.common import MockConfigEntry, async_fire_time_changed
 
 
 async def test_ip_changes_fallback_discovery(hass: HomeAssistant):
@@ -51,19 +54,15 @@ async def test_ip_changes_fallback_discovery(hass: HomeAssistant):
 
     mocked_bulb = _mocked_bulb(True)
     mocked_bulb.bulb_type = BulbType.WhiteTempMood
-    mocked_bulb.get_capabilities = MagicMock(
+    mocked_bulb.async_listen = AsyncMock(side_effect=[BulbException, None, None, None])
-        side_effect=[OSError, CAPABILITIES, CAPABILITIES]
-    )
 
-    _discovered_devices = [{"capabilities": CAPABILITIES, "ip": IP_ADDRESS}]
+    with patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb), _patch_discovery():
-    with patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb), patch(
-        f"{MODULE}.discover_bulbs", return_value=_discovered_devices
-    ):
         assert await hass.config_entries.async_setup(config_entry.entry_id)
         await hass.async_block_till_done()
+        await hass.async_block_till_done()
 
     binary_sensor_entity_id = ENTITY_BINARY_SENSOR_TEMPLATE.format(
-        f"yeelight_color_{ID}"
+        f"yeelight_color_{SHORT_ID}"
     )
 
     type(mocked_bulb).async_get_properties = AsyncMock(None)
@@ -77,6 +76,19 @@ async def test_ip_changes_fallback_discovery(hass: HomeAssistant):
     entity_registry = er.async_get(hass)
     assert entity_registry.async_get(binary_sensor_entity_id) is not None
 
+    with patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb), _patch_discovery():
+        # The discovery should update the ip address
+        async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=2))
+        await hass.async_block_till_done()
+        assert config_entry.data[CONF_HOST] == IP_ADDRESS
+
+    # Make sure we can still reload with the new ip right after we change it
+    with patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb), _patch_discovery():
+        await hass.config_entries.async_reload(config_entry.entry_id)
+        await hass.async_block_till_done()
+
+    assert entity_registry.async_get(binary_sensor_entity_id) is not None
+
 
 async def test_ip_changes_id_missing_cannot_fallback(hass: HomeAssistant):
     """Test Yeelight ip changes and we fallback to discovery."""
@@ -85,9 +97,7 @@ async def test_ip_changes_id_missing_cannot_fallback(hass: HomeAssistant):
 
     mocked_bulb = _mocked_bulb(True)
     mocked_bulb.bulb_type = BulbType.WhiteTempMood
-    mocked_bulb.get_capabilities = MagicMock(
+    mocked_bulb.async_listen = AsyncMock(side_effect=[BulbException, None, None, None])
-        side_effect=[OSError, CAPABILITIES, CAPABILITIES]
-    )
 
     with patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb):
         assert not await hass.config_entries.async_setup(config_entry.entry_id)
@@ -102,9 +112,7 @@ async def test_setup_discovery(hass: HomeAssistant):
     config_entry.add_to_hass(hass)
 
     mocked_bulb = _mocked_bulb()
-    with _patch_discovery(MODULE), patch(
+    with _patch_discovery(), patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb):
-        f"{MODULE}.AsyncBulb", return_value=mocked_bulb
-    ):
         assert await hass.config_entries.async_setup(config_entry.entry_id)
         await hass.async_block_till_done()
 
@@ -127,9 +135,7 @@ async def test_setup_import(hass: HomeAssistant):
     """Test import from yaml."""
     mocked_bulb = _mocked_bulb()
     name = "yeelight"
-    with patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb), patch(
+    with patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb), _patch_discovery():
-        f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb
-    ):
         assert await async_setup_component(
             hass,
             DOMAIN,
@@ -162,9 +168,7 @@ async def test_unique_ids_device(hass: HomeAssistant):
 
     mocked_bulb = _mocked_bulb()
     mocked_bulb.bulb_type = BulbType.WhiteTempMood
-    with _patch_discovery(MODULE), patch(
+    with _patch_discovery(), patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb):
-        f"{MODULE}.AsyncBulb", return_value=mocked_bulb
-    ):
         assert await hass.config_entries.async_setup(config_entry.entry_id)
         await hass.async_block_till_done()
 
@@ -188,9 +192,7 @@ async def test_unique_ids_entry(hass: HomeAssistant):
     mocked_bulb = _mocked_bulb()
     mocked_bulb.bulb_type = BulbType.WhiteTempMood
 
-    with _patch_discovery(MODULE), patch(
+    with _patch_discovery(), patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb):
-        f"{MODULE}.AsyncBulb", return_value=mocked_bulb
-    ):
         assert await hass.config_entries.async_setup(config_entry.entry_id)
         await hass.async_block_till_done()
 
@@ -220,30 +222,13 @@ async def test_bulb_off_while_adding_in_ha(hass: HomeAssistant):
     mocked_bulb = _mocked_bulb(True)
     mocked_bulb.bulb_type = BulbType.WhiteTempMood
 
-    with patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb), patch(
+    with patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb), _patch_discovery(
-        f"{MODULE}.config_flow.yeelight.Bulb", return_value=mocked_bulb
+        no_device=True
-    ):
+    ), _patch_discovery_timeout(), _patch_discovery_interval():
         assert await hass.config_entries.async_setup(config_entry.entry_id)
         await hass.async_block_till_done()
 
-    binary_sensor_entity_id = ENTITY_BINARY_SENSOR_TEMPLATE.format(
+    assert config_entry.state is ConfigEntryState.LOADED
-        IP_ADDRESS.replace(".", "_")
-    )
-
-    type(mocked_bulb).get_capabilities = MagicMock(CAPABILITIES)
-    type(mocked_bulb).get_properties = MagicMock(None)
-
-    await hass.data[DOMAIN][DATA_CONFIG_ENTRIES][config_entry.entry_id][
-        DATA_DEVICE
-    ].async_update()
-    hass.data[DOMAIN][DATA_CONFIG_ENTRIES][config_entry.entry_id][
-        DATA_DEVICE
-    ].async_update_callback({})
-    await hass.async_block_till_done()
-    await hass.async_block_till_done()
-
-    entity_registry = er.async_get(hass)
-    assert entity_registry.async_get(binary_sensor_entity_id) is not None
 
 
 async def test_async_listen_error_late_discovery(hass, caplog):
@@ -251,12 +236,9 @@ async def test_async_listen_error_late_discovery(hass, caplog):
     config_entry = MockConfigEntry(domain=DOMAIN, data=CONFIG_ENTRY_DATA)
     config_entry.add_to_hass(hass)
 
-    mocked_bulb = _mocked_bulb()
+    mocked_bulb = _mocked_bulb(cannot_connect=True)
-    mocked_bulb.async_listen = AsyncMock(side_effect=BulbException)
 
-    with _patch_discovery(MODULE), patch(
+    with _patch_discovery(), patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb):
-        f"{MODULE}.AsyncBulb", return_value=mocked_bulb
-    ):
         await hass.config_entries.async_setup(config_entry.entry_id)
         await hass.async_block_till_done()
 
@@ -264,17 +246,33 @@ async def test_async_listen_error_late_discovery(hass, caplog):
     assert "Failed to connect to bulb at" in caplog.text
 
 
-async def test_async_listen_error_has_host(hass: HomeAssistant):
+async def test_async_listen_error_has_host_with_id(hass: HomeAssistant):
     """Test the async listen error."""
     config_entry = MockConfigEntry(
         domain=DOMAIN, data={CONF_ID: ID, CONF_HOST: "127.0.0.1"}
     )
     config_entry.add_to_hass(hass)
 
-    mocked_bulb = _mocked_bulb()
+    with _patch_discovery(
-    mocked_bulb.async_listen = AsyncMock(side_effect=BulbException)
+        no_device=True
+    ), _patch_discovery_timeout(), _patch_discovery_interval(), patch(
+        f"{MODULE}.AsyncBulb", return_value=_mocked_bulb(cannot_connect=True)
+    ):
+        await hass.config_entries.async_setup(config_entry.entry_id)
 
-    with patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb):
+    assert config_entry.state is ConfigEntryState.LOADED
+
+
+async def test_async_listen_error_has_host_without_id(hass: HomeAssistant):
+    """Test the async listen error but no id."""
+    config_entry = MockConfigEntry(domain=DOMAIN, data={CONF_HOST: "127.0.0.1"})
+    config_entry.add_to_hass(hass)
+
+    with _patch_discovery(
+        no_device=True
+    ), _patch_discovery_timeout(), _patch_discovery_interval(), patch(
+        f"{MODULE}.AsyncBulb", return_value=_mocked_bulb(cannot_connect=True)
+    ):
         await hass.config_entries.async_setup(config_entry.entry_id)
 
     assert config_entry.state is ConfigEntryState.SETUP_RETRY

tests/components/yeelight/test_light.py

@@ -102,9 +102,10 @@ from . import (
     MODULE,
     NAME,
     PROPERTIES,
-    UNIQUE_NAME,
+    UNIQUE_FRIENDLY_NAME,
     _mocked_bulb,
     _patch_discovery,
+    _patch_discovery_interval,
 )
 
 from tests.common import MockConfigEntry
@@ -132,7 +133,7 @@ async def test_services(hass: HomeAssistant, caplog):
     config_entry.add_to_hass(hass)
 
     mocked_bulb = _mocked_bulb()
-    with _patch_discovery(MODULE), patch(
+    with _patch_discovery(), _patch_discovery_interval(), patch(
         f"{MODULE}.AsyncBulb", return_value=mocked_bulb
     ):
         assert await hass.config_entries.async_setup(config_entry.entry_id)
@@ -559,7 +560,7 @@ async def test_device_types(hass: HomeAssistant, caplog):
         model,
         target_properties,
         nightlight_properties=None,
-        name=UNIQUE_NAME,
+        name=UNIQUE_FRIENDLY_NAME,
         entity_id=ENTITY_LIGHT,
     ):
         config_entry = MockConfigEntry(
@@ -598,7 +599,7 @@ async def test_device_types(hass: HomeAssistant, caplog):
         assert hass.states.get(entity_id).state == "off"
         state = hass.states.get(f"{entity_id}_nightlight")
         assert state.state == "on"
-        nightlight_properties["friendly_name"] = f"{name} nightlight"
+        nightlight_properties["friendly_name"] = f"{name} Nightlight"
         nightlight_properties["icon"] = "mdi:weather-night"
         nightlight_properties["flowing"] = False
         nightlight_properties["night_light"] = True
@@ -893,7 +894,7 @@ async def test_device_types(hass: HomeAssistant, caplog):
             "color_mode": "color_temp",
             "supported_color_modes": ["color_temp", "hs", "rgb"],
         },
-        name=f"{UNIQUE_NAME} ambilight",
+        name=f"{UNIQUE_FRIENDLY_NAME} Ambilight",
         entity_id=f"{ENTITY_LIGHT}_ambilight",
     )
 
@@ -914,7 +915,7 @@ async def test_device_types(hass: HomeAssistant, caplog):
             "color_mode": "hs",
             "supported_color_modes": ["color_temp", "hs", "rgb"],
         },
-        name=f"{UNIQUE_NAME} ambilight",
+        name=f"{UNIQUE_FRIENDLY_NAME} Ambilight",
         entity_id=f"{ENTITY_LIGHT}_ambilight",
     )
 
@@ -935,7 +936,7 @@ async def test_device_types(hass: HomeAssistant, caplog):
             "color_mode": "rgb",
             "supported_color_modes": ["color_temp", "hs", "rgb"],
         },
-        name=f"{UNIQUE_NAME} ambilight",
+        name=f"{UNIQUE_FRIENDLY_NAME} Ambilight",
         entity_id=f"{ENTITY_LIGHT}_ambilight",
     )
 
@@ -969,7 +970,7 @@ async def test_effects(hass: HomeAssistant):
     config_entry.add_to_hass(hass)
 
     mocked_bulb = _mocked_bulb()
-    with _patch_discovery(MODULE), patch(
+    with _patch_discovery(), _patch_discovery_interval(), patch(
         f"{MODULE}.AsyncBulb", return_value=mocked_bulb
     ):
         assert await hass.config_entries.async_setup(config_entry.entry_id)