From 849cfa3af818867c5d377e9e1337fb7d977a1468 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Thu, 24 Aug 2023 05:04:00 -0500
Subject: [PATCH] Retry yeelight setup later if the wrong device is found
 (#98884)

---
 homeassistant/components/yeelight/__init__.py | 14 +++++++++++
 homeassistant/components/yeelight/device.py   | 20 ++++++++++-----
 tests/components/yeelight/test_init.py        | 25 +++++++++++++++++++
 3 files changed, 53 insertions(+), 6 deletions(-)

diff --git a/homeassistant/components/yeelight/__init__.py b/homeassistant/components/yeelight/__init__.py
index c07852629a9..cc9faa33194 100644
--- a/homeassistant/components/yeelight/__init__.py
+++ b/homeassistant/components/yeelight/__init__.py
@@ -144,6 +144,7 @@ async def _async_initialize(
     entry: ConfigEntry,
     device: YeelightDevice,
 ) -> None:
+    """Initialize a Yeelight device."""
     entry_data = hass.data[DOMAIN][DATA_CONFIG_ENTRIES][entry.entry_id] = {}
     await device.async_setup()
     entry_data[DATA_DEVICE] = device
@@ -216,6 +217,19 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     except (asyncio.TimeoutError, OSError, BulbException) as ex:
         raise ConfigEntryNotReady from ex
 
+    found_unique_id = device.unique_id
+    expected_unique_id = entry.unique_id
+    if expected_unique_id and found_unique_id and found_unique_id != expected_unique_id:
+        # If the id of the device does not match the unique_id
+        # of the config entry, it likely means the DHCP lease has expired
+        # and the device has been assigned a new IP address. We need to
+        # wait for the next discovery to find the device at its new address
+        # and update the config entry so we do not mix up devices.
+        raise ConfigEntryNotReady(
+            f"Unexpected device found at {device.host}; "
+            f"expected {expected_unique_id}, found {found_unique_id}"
+        )
+
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     # Wait to install the reload listener until everything was successfully initialized
diff --git a/homeassistant/components/yeelight/device.py b/homeassistant/components/yeelight/device.py
index 0fabe693aa9..811a1904b04 100644
--- a/homeassistant/components/yeelight/device.py
+++ b/homeassistant/components/yeelight/device.py
@@ -3,12 +3,13 @@ from __future__ import annotations
 
 import asyncio
 import logging
+from typing import Any
 
 from yeelight import BulbException
-from yeelight.aio import KEY_CONNECTED
+from yeelight.aio import KEY_CONNECTED, AsyncBulb
 
 from homeassistant.const import CONF_ID, CONF_NAME
-from homeassistant.core import callback
+from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_send
 from homeassistant.helpers.event import async_call_later
 
@@ -63,17 +64,19 @@ def update_needs_bg_power_workaround(data):
 class YeelightDevice:
     """Represents single Yeelight device."""
 
-    def __init__(self, hass, host, config, bulb):
+    def __init__(
+        self, hass: HomeAssistant, host: str, config: dict[str, Any], bulb: AsyncBulb
+    ) -> None:
         """Initialize device."""
         self._hass = hass
         self._config = config
         self._host = host
         self._bulb_device = bulb
-        self.capabilities = {}
-        self._device_type = None
+        self.capabilities: dict[str, Any] = {}
+        self._device_type: str | None = None
         self._available = True
         self._initialized = False
-        self._name = None
+        self._name: str | None = None
 
     @property
     def bulb(self):
@@ -115,6 +118,11 @@ class YeelightDevice:
         """Return the firmware version."""
         return self.capabilities.get("fw_ver")
 
+    @property
+    def unique_id(self) -> str | None:
+        """Return the unique ID of the device."""
+        return self.capabilities.get("id")
+
     @property
     def is_nightlight_supported(self) -> bool:
         """Return true / false if nightlight is supported.
diff --git a/tests/components/yeelight/test_init.py b/tests/components/yeelight/test_init.py
index 906dbf50ace..b439ce04c25 100644
--- a/tests/components/yeelight/test_init.py
+++ b/tests/components/yeelight/test_init.py
@@ -618,3 +618,28 @@ async def test_async_setup_with_discovery_not_working(hass: HomeAssistant) -> No
         assert config_entry.state is ConfigEntryState.LOADED
 
     assert hass.states.get("light.yeelight_color_0x15243f").state == STATE_ON
+
+
+async def test_async_setup_retries_with_wrong_device(
+    hass: HomeAssistant, caplog: pytest.LogCaptureFixture
+) -> None:
+    """Test the config entry enters a retry state with the wrong device."""
+    config_entry = MockConfigEntry(
+        domain=DOMAIN,
+        data={CONF_HOST: IP_ADDRESS, CONF_ID: "0x0000000000999999"},
+        options={},
+        unique_id="0x0000000000999999",
+    )
+    config_entry.add_to_hass(hass)
+
+    with _patch_discovery(), _patch_discovery_timeout(), _patch_discovery_interval(), patch(
+        f"{MODULE}.AsyncBulb", return_value=_mocked_bulb()
+    ):
+        await hass.config_entries.async_setup(config_entry.entry_id)
+        await hass.async_block_till_done()
+
+    assert config_entry.state is ConfigEntryState.SETUP_RETRY
+    assert (
+        "Unexpected device found at 192.168.1.239; expected 0x0000000000999999, "
+        "found 0x000000000015243f; Retrying in background"
+    ) in caplog.text