Add discovery to yeelight (#50385)

homeassistant/components/yeelight/__init__.py

@@ -8,7 +8,7 @@ import logging
 import voluptuous as vol
 from yeelight import Bulb, BulbException, discover_bulbs
 
-from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry
+from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry, ConfigEntryNotReady
 from homeassistant.const import (
     CONF_DEVICES,
     CONF_HOST,
@@ -48,8 +48,8 @@ DATA_CONFIG_ENTRIES = "config_entries"
 DATA_CUSTOM_EFFECTS = "custom_effects"
 DATA_SCAN_INTERVAL = "scan_interval"
 DATA_DEVICE = "device"
-DATA_UNSUB_UPDATE_LISTENER = "unsub_update_listener"
 DATA_REMOVE_INIT_DISPATCHER = "remove_init_dispatcher"
+DATA_PLATFORMS_LOADED = "platforms_loaded"
 
 ATTR_COUNT = "count"
 ATTR_ACTION = "action"
@@ -179,81 +179,115 @@ async def async_setup(hass: HomeAssistant, config: dict) -> bool:
     return True
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def _async_initialize(
-    """Set up Yeelight from a config entry."""
+    hass: HomeAssistant,
+    entry: ConfigEntry,
+    host: str,
+    device: YeelightDevice | None = None,
+) -> None:
+    entry_data = hass.data[DOMAIN][DATA_CONFIG_ENTRIES][entry.entry_id] = {
+        DATA_PLATFORMS_LOADED: False
+    }
+    entry.async_on_unload(entry.add_update_listener(_async_update_listener))
 
-    async def _initialize(host: str, capabilities: dict | None = None) -> None:
+    @callback
-        remove_dispatcher = async_dispatcher_connect(
+    def _async_load_platforms():
-            hass,
+        if entry_data[DATA_PLATFORMS_LOADED]:
-            DEVICE_INITIALIZED.format(host),
+            return
-            _load_platforms,
+        entry_data[DATA_PLATFORMS_LOADED] = True
-        )
-        hass.data[DOMAIN][DATA_CONFIG_ENTRIES][entry.entry_id][
-            DATA_REMOVE_INIT_DISPATCHER
-        ] = remove_dispatcher
-
-        device = await _async_get_device(hass, host, entry, capabilities)
-        hass.data[DOMAIN][DATA_CONFIG_ENTRIES][entry.entry_id][DATA_DEVICE] = device
-
-        await device.async_setup()
-
-    async def _load_platforms():
         hass.config_entries.async_setup_platforms(entry, PLATFORMS)
 
-    # Move options from data for imported entries
+    if not device:
-    # Initialize options with default values for other entries
+        device = await _async_get_device(hass, host, entry)
-    if not entry.options:
+    entry_data[DATA_DEVICE] = device
-        hass.config_entries.async_update_entry(
-            entry,
-            data={
-                CONF_HOST: entry.data.get(CONF_HOST),
-                CONF_ID: entry.data.get(CONF_ID),
-            },
-            options={
-                CONF_NAME: entry.data.get(CONF_NAME, ""),
-                CONF_MODEL: entry.data.get(CONF_MODEL, ""),
-                CONF_TRANSITION: entry.data.get(CONF_TRANSITION, DEFAULT_TRANSITION),
-                CONF_MODE_MUSIC: entry.data.get(CONF_MODE_MUSIC, DEFAULT_MODE_MUSIC),
-                CONF_SAVE_ON_CHANGE: entry.data.get(
-                    CONF_SAVE_ON_CHANGE, DEFAULT_SAVE_ON_CHANGE
-                ),
-                CONF_NIGHTLIGHT_SWITCH: entry.data.get(
-                    CONF_NIGHTLIGHT_SWITCH, DEFAULT_NIGHTLIGHT_SWITCH
-                ),
-            },
-        )
 
-    hass.data[DOMAIN][DATA_CONFIG_ENTRIES][entry.entry_id] = {
+    entry.async_on_unload(
-        DATA_UNSUB_UPDATE_LISTENER: entry.add_update_listener(_async_update_listener)
+        async_dispatcher_connect(
-    }
+            hass,
+            DEVICE_INITIALIZED.format(host),
+            _async_load_platforms,
+        )
+    )
+
+    entry.async_on_unload(device.async_unload)
+    await device.async_setup()
+
+
+@callback
+def _async_populate_entry_options(hass: HomeAssistant, entry: ConfigEntry) -> None:
+    """Move options from data for imported entries.
+
+    Initialize options with default values for other entries.
+    """
+    if entry.options:
+        return
+
+    hass.config_entries.async_update_entry(
+        entry,
+        data={
+            CONF_HOST: entry.data.get(CONF_HOST),
+            CONF_ID: entry.data.get(CONF_ID),
+        },
+        options={
+            CONF_NAME: entry.data.get(CONF_NAME, ""),
+            CONF_MODEL: entry.data.get(CONF_MODEL, ""),
+            CONF_TRANSITION: entry.data.get(CONF_TRANSITION, DEFAULT_TRANSITION),
+            CONF_MODE_MUSIC: entry.data.get(CONF_MODE_MUSIC, DEFAULT_MODE_MUSIC),
+            CONF_SAVE_ON_CHANGE: entry.data.get(
+                CONF_SAVE_ON_CHANGE, DEFAULT_SAVE_ON_CHANGE
+            ),
+            CONF_NIGHTLIGHT_SWITCH: entry.data.get(
+                CONF_NIGHTLIGHT_SWITCH, DEFAULT_NIGHTLIGHT_SWITCH
+            ),
+        },
+    )
+
+
+async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+    """Set up Yeelight from a config entry."""
+    _async_populate_entry_options(hass, entry)
 
     if entry.data.get(CONF_HOST):
-        # manually added device
+        try:
-        await _initialize(entry.data[CONF_HOST])
+            device = await _async_get_device(hass, entry.data[CONF_HOST], entry)
-    else:
+        except OSError as ex:
-        # discovery
+            # If CONF_ID is not valid we cannot fallback to discovery
-        scanner = YeelightScanner.async_get(hass)
+            # so we must retry by raising ConfigEntryNotReady
-        scanner.async_register_callback(entry.data[CONF_ID], _initialize)
+            if not entry.data.get(CONF_ID):
+                raise ConfigEntryNotReady from ex
+            # Otherwise fall through to discovery
+        else:
+            # manually added device
+            await _async_initialize(hass, entry, entry.data[CONF_HOST], device=device)
+            return True
 
+    # discovery
+    scanner = YeelightScanner.async_get(hass)
+
+    async def _async_from_discovery(host: str) -> None:
+        await _async_initialize(hass, entry, host)
+
+    scanner.async_register_callback(entry.data[CONF_ID], _async_from_discovery)
     return True
 
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry):
     """Unload a config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
+    data_config_entries = hass.data[DOMAIN][DATA_CONFIG_ENTRIES]
-    if unload_ok:
+    entry_data = data_config_entries[entry.entry_id]
-        data = hass.data[DOMAIN][DATA_CONFIG_ENTRIES].pop(entry.entry_id)
-        remove_init_dispatcher = data.get(DATA_REMOVE_INIT_DISPATCHER)
-        if remove_init_dispatcher is not None:
-            remove_init_dispatcher()
-        data[DATA_UNSUB_UPDATE_LISTENER]()
-        data[DATA_DEVICE].async_unload()
-        if entry.data[CONF_ID]:
-            # discovery
-            scanner = YeelightScanner.async_get(hass)
-            scanner.async_unregister_callback(entry.data[CONF_ID])
 
-    return unload_ok
+    if entry_data[DATA_PLATFORMS_LOADED]:
+        if not await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
+            return False
+
+    if entry.data.get(CONF_ID):
+        # discovery
+        scanner = YeelightScanner.async_get(hass)
+        scanner.async_unregister_callback(entry.data[CONF_ID])
+
+    data_config_entries.pop(entry.entry_id)
+
+    return True
 
 
 @callback
@@ -582,16 +616,12 @@ async def _async_get_device(
     hass: HomeAssistant,
     host: str,
     entry: ConfigEntry,
-    capabilities: dict | None,
 ) -> YeelightDevice:
     # Get model from config and capabilities
     model = entry.options.get(CONF_MODEL)
-    if not model and capabilities is not None:
-        model = capabilities.get("model")
 
     # Set up device
     bulb = Bulb(host, model=model or None)
-    if capabilities is None:
+    capabilities = await hass.async_add_executor_job(bulb.get_capabilities)
-        capabilities = await hass.async_add_executor_job(bulb.get_capabilities)
 
     return YeelightDevice(hass, host, entry.options, bulb, capabilities)

homeassistant/components/yeelight/config_flow.py

@@ -5,6 +5,7 @@ import voluptuous as vol
 import yeelight
 
 from homeassistant import config_entries, exceptions
+from homeassistant.components.dhcp import IP_ADDRESS
 from homeassistant.const import CONF_DEVICE, CONF_HOST, CONF_ID, CONF_NAME
 from homeassistant.core import callback
 import homeassistant.helpers.config_validation as cv
@@ -21,6 +22,8 @@ from . import (
     _async_unique_name,
 )
 
+MODEL_UNKNOWN = "unknown"
+
 _LOGGER = logging.getLogger(__name__)
 
 
@@ -38,22 +41,69 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
     def __init__(self):
         """Initialize the config flow."""
         self._discovered_devices = {}
+        self._discovered_model = None
+        self._discovered_ip = None
+
+    async def async_step_homekit(self, discovery_info):
+        """Handle discovery from homekit."""
+        self._discovered_ip = discovery_info["host"]
+        return await self._async_handle_discovery()
+
+    async def async_step_dhcp(self, discovery_info):
+        """Handle discovery from dhcp."""
+        self._discovered_ip = discovery_info[IP_ADDRESS]
+        return await self._async_handle_discovery()
+
+    async def _async_handle_discovery(self):
+        """Handle any discovery."""
+        self.context[CONF_HOST] = self._discovered_ip
+        for progress in self._async_in_progress():
+            if progress.get("context", {}).get(CONF_HOST) == self._discovered_ip:
+                return self.async_abort(reason="already_in_progress")
+
+        self._discovered_model = await self._async_try_connect(self._discovered_ip)
+        if not self.unique_id:
+            return self.async_abort(reason="cannot_connect")
+
+        self._abort_if_unique_id_configured(
+            updates={CONF_HOST: self._discovered_ip}, reload_on_update=False
+        )
+        return await self.async_step_discovery_confirm()
+
+    async def async_step_discovery_confirm(self, user_input=None):
+        """Confirm discovery."""
+        if user_input is not None:
+            return self.async_create_entry(
+                title=f"{self._discovered_model} {self.unique_id}",
+                data={CONF_ID: self.unique_id, CONF_HOST: self._discovered_ip},
+            )
+
+        self._set_confirm_only()
+        placeholders = {
+            "model": self._discovered_model,
+            "host": self._discovered_ip,
+        }
+        self.context["title_placeholders"] = placeholders
+        return self.async_show_form(
+            step_id="discovery_confirm", description_placeholders=placeholders
+        )
 
     async def async_step_user(self, user_input=None):
         """Handle the initial step."""
         errors = {}
         if user_input is not None:
-            if user_input.get(CONF_HOST):
+            if not user_input.get(CONF_HOST):
-                try:
-                    await self._async_try_connect(user_input[CONF_HOST])
-                    return self.async_create_entry(
-                        title=user_input[CONF_HOST],
-                        data=user_input,
-                    )
-                except CannotConnect:
-                    errors["base"] = "cannot_connect"
-            else:
                 return await self.async_step_pick_device()
+            try:
+                model = await self._async_try_connect(user_input[CONF_HOST])
+            except CannotConnect:
+                errors["base"] = "cannot_connect"
+            else:
+                self._abort_if_unique_id_configured()
+                return self.async_create_entry(
+                    title=f"{model} {self.unique_id}",
+                    data=user_input,
+                )
 
         user_input = user_input or {}
         return self.async_show_form(
@@ -117,6 +167,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
                 user_input.pop(CONF_NIGHTLIGHT_SWITCH_TYPE)
                 == NIGHTLIGHT_SWITCH_TYPE_LIGHT
             )
+        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):
@@ -131,8 +182,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
             else:
                 _LOGGER.debug("Get capabilities: %s", capabilities)
                 await self.async_set_unique_id(capabilities["id"])
-                self._abort_if_unique_id_configured()
+                return capabilities["model"]
-                return
         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
@@ -145,6 +195,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
             _LOGGER.error("Failed to get properties from %s: %s", host, err)
             raise CannotConnect from err
         _LOGGER.debug("Get properties: %s", bulb.last_properties)
+        return MODEL_UNKNOWN
 
 
 class OptionsFlowHandler(config_entries.OptionsFlow):

homeassistant/components/yeelight/manifest.json

@@ -5,5 +5,11 @@
   "requirements": ["yeelight==0.6.2"],
   "codeowners": ["@rytilahti", "@zewelor", "@shenxn"],
   "config_flow": true,
-  "iot_class": "local_polling"
+  "iot_class": "local_polling",
+  "dhcp": [{
+    "hostname": "yeelink-*"
+  }],
+  "homekit": {
+    "models": ["YLDP*"]
+  }
 }

homeassistant/components/yeelight/strings.json

@@ -1,5 +1,6 @@
 {
   "config": {
+    "flow_title": "{model} {host}",
     "step": {
       "user": {
         "description": "If you leave the host empty, discovery will be used to find devices.",
@@ -11,6 +12,9 @@
         "data": {
           "device": "Device"
         }
+      },
+      "discovery_confirm": {
+        "description": "Do you want to setup {model} ({host})?"
       }
     },
     "error": {

homeassistant/components/yeelight/translations/en.json

@@ -7,7 +7,11 @@
         "error": {
             "cannot_connect": "Failed to connect"
         },
+        "flow_title": "{model} {host}",
         "step": {
+            "discovery_confirm": {
+                "description": "Do you want to setup {model} ({host})?"
+            },
             "pick_device": {
                 "data": {
                     "device": "Device"

homeassistant/generated/dhcp.py

@@ -343,5 +343,9 @@ DHCP = [
     {
         "domain": "verisure",
         "macaddress": "0023C1*"
+    },
+    {
+        "domain": "yeelight",
+        "hostname": "yeelink-*"
     }
 ]

homeassistant/generated/zeroconf.py

@@ -240,6 +240,7 @@ HOMEKIT = {
     "Touch HD": "rainmachine",
     "Welcome": "netatmo",
     "Wemo": "wemo",
+    "YLDP*": "yeelight",
     "iSmartGate": "gogogate2",
     "iZone": "izone",
     "tado": "tado"

tests/components/yeelight/test_config_flow.py

@@ -1,7 +1,9 @@
 """Test the Yeelight config flow."""
 from unittest.mock import MagicMock, patch
 
-from homeassistant import config_entries
+import pytest
+
+from homeassistant import config_entries, setup
 from homeassistant.components.yeelight import (
     CONF_MODE_MUSIC,
     CONF_MODEL,
@@ -19,6 +21,7 @@ from homeassistant.components.yeelight import (
 )
 from homeassistant.const import CONF_DEVICE, CONF_HOST, CONF_ID, CONF_NAME
 from homeassistant.core import HomeAssistant
+from homeassistant.data_entry_flow import RESULT_TYPE_ABORT, RESULT_TYPE_FORM
 
 from . import (
     ID,
@@ -205,7 +208,7 @@ async def test_manual(hass: HomeAssistant):
         )
         await hass.async_block_till_done()
     assert result4["type"] == "create_entry"
-    assert result4["title"] == IP_ADDRESS
+    assert result4["title"] == "color 0x000000000015243f"
     assert result4["data"] == {CONF_HOST: IP_ADDRESS}
 
     # Duplicate
@@ -286,3 +289,103 @@ async def test_manual_no_capabilities(hass: HomeAssistant):
     type(mocked_bulb).get_properties.assert_called_once()
     assert result["type"] == "create_entry"
     assert result["data"] == {CONF_HOST: IP_ADDRESS}
+
+
+async def test_discovered_by_homekit_and_dhcp(hass):
+    """Test we get the form with homekit and abort for dhcp source when we get both."""
+    await setup.async_setup_component(hass, "persistent_notification", {})
+
+    mocked_bulb = _mocked_bulb()
+    with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", 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"}},
+        )
+    assert result["type"] == RESULT_TYPE_FORM
+    assert result["errors"] is None
+
+    with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", 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"},
+        )
+    assert result2["type"] == RESULT_TYPE_ABORT
+    assert result2["reason"] == "already_in_progress"
+
+    with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", 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"},
+        )
+    assert result3["type"] == RESULT_TYPE_ABORT
+    assert result3["reason"] == "already_in_progress"
+
+
+@pytest.mark.parametrize(
+    "source, data",
+    [
+        (
+            config_entries.SOURCE_DHCP,
+            {"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"}},
+        ),
+    ],
+)
+async def test_discovered_by_dhcp_or_homekit(hass, source, data):
+    """Test we can setup when discovered from dhcp or homekit."""
+    await setup.async_setup_component(hass, "persistent_notification", {})
+
+    mocked_bulb = _mocked_bulb()
+    with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb):
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN,
+            context={"source": source},
+            data=data,
+        )
+    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(
+        f"{MODULE}.async_setup_entry",
+        return_value=True,
+    ) as mock_async_setup_entry:
+        result2 = await hass.config_entries.flow.async_configure(result["flow_id"], {})
+    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
+
+
+@pytest.mark.parametrize(
+    "source, data",
+    [
+        (
+            config_entries.SOURCE_DHCP,
+            {"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"}},
+        ),
+    ],
+)
+async def test_discovered_by_dhcp_or_homekit_failed_to_get_id(hass, source, data):
+    """Test we abort if we cannot get the unique id when discovered from dhcp or homekit."""
+    await setup.async_setup_component(hass, "persistent_notification", {})
+
+    mocked_bulb = _mocked_bulb()
+    type(mocked_bulb).get_capabilities = MagicMock(return_value=None)
+    with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", 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"

tests/components/yeelight/test_init.py

@@ -11,7 +11,14 @@ from homeassistant.components.yeelight import (
     DOMAIN,
     NIGHTLIGHT_SWITCH_TYPE_LIGHT,
 )
-from homeassistant.const import CONF_DEVICES, CONF_HOST, CONF_NAME, STATE_UNAVAILABLE
+from homeassistant.config_entries import ENTRY_STATE_SETUP_RETRY
+from homeassistant.const import (
+    CONF_DEVICES,
+    CONF_HOST,
+    CONF_ID,
+    CONF_NAME,
+    STATE_UNAVAILABLE,
+)
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
 from homeassistant.setup import async_setup_component
@@ -35,6 +42,77 @@ from . import (
 from tests.common import MockConfigEntry
 
 
+async def test_ip_changes_fallback_discovery(hass: HomeAssistant):
+    """Test Yeelight ip changes and we fallback to discovery."""
+    config_entry = MockConfigEntry(
+        domain=DOMAIN,
+        data={
+            CONF_ID: ID,
+            CONF_HOST: "5.5.5.5",
+        },
+        unique_id=ID,
+    )
+    config_entry.add_to_hass(hass)
+
+    mocked_bulb = _mocked_bulb(True)
+    mocked_bulb.bulb_type = BulbType.WhiteTempMood
+    mocked_bulb.get_capabilities = MagicMock(
+        side_effect=[OSError, CAPABILITIES, CAPABILITIES]
+    )
+
+    _discovered_devices = [
+        {
+            "capabilities": CAPABILITIES,
+            "ip": IP_ADDRESS,
+        }
+    ]
+    with patch(f"{MODULE}.Bulb", 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()
+
+    binary_sensor_entity_id = ENTITY_BINARY_SENSOR_TEMPLATE.format(
+        f"yeelight_color_{ID}"
+    )
+    entity_registry = er.async_get(hass)
+    assert entity_registry.async_get(binary_sensor_entity_id) is None
+
+    await hass.async_block_till_done()
+
+    type(mocked_bulb).get_properties = MagicMock(None)
+
+    hass.data[DOMAIN][DATA_CONFIG_ENTRIES][config_entry.entry_id][DATA_DEVICE].update()
+    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_ip_changes_id_missing_cannot_fallback(hass: HomeAssistant):
+    """Test Yeelight ip changes and we fallback to discovery."""
+    config_entry = MockConfigEntry(
+        domain=DOMAIN,
+        data={
+            CONF_HOST: "5.5.5.5",
+        },
+    )
+    config_entry.add_to_hass(hass)
+
+    mocked_bulb = _mocked_bulb(True)
+    mocked_bulb.bulb_type = BulbType.WhiteTempMood
+    mocked_bulb.get_capabilities = MagicMock(
+        side_effect=[OSError, CAPABILITIES, CAPABILITIES]
+    )
+
+    with patch(f"{MODULE}.Bulb", return_value=mocked_bulb):
+        assert not await hass.config_entries.async_setup(config_entry.entry_id)
+        await hass.async_block_till_done()
+
+    assert config_entry.state == ENTRY_STATE_SETUP_RETRY
+
+
 async def test_setup_discovery(hass: HomeAssistant):
     """Test setting up Yeelight by discovery."""
     config_entry = MockConfigEntry(domain=DOMAIN, data=CONFIG_ENTRY_DATA)
@@ -182,6 +260,7 @@ async def test_bulb_off_while_adding_in_ha(hass: HomeAssistant):
 
     hass.data[DOMAIN][DATA_CONFIG_ENTRIES][config_entry.entry_id][DATA_DEVICE].update()
     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