Fix yeelight config flow ip update and timeout (#66883)

homeassistant/components/yeelight/config_flow.py

@@ -1,4 +1,5 @@
 """Config flow for Yeelight integration."""
+import asyncio
 import logging
 from urllib.parse import urlparse
 
@@ -86,11 +87,13 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
 
     async def _async_handle_discovery_with_unique_id(self):
         """Handle any discovery with a unique id."""
-        for entry in self._async_current_entries():
-            if entry.unique_id != self.unique_id:
+        for entry in self._async_current_entries(include_ignore=False):
+            if entry.unique_id != self.unique_id and self.unique_id != entry.data.get(
+                CONF_ID
+            ):
                 continue
             reload = entry.state == ConfigEntryState.SETUP_RETRY
-            if entry.data[CONF_HOST] != self._discovered_ip:
+            if entry.data.get(CONF_HOST) != self._discovered_ip:
                 self.hass.config_entries.async_update_entry(
                     entry, data={**entry.data, CONF_HOST: self._discovered_ip}
                 )
@@ -261,7 +264,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
             await bulb.async_listen(lambda _: True)
             await bulb.async_get_properties()
             await bulb.async_stop_listening()
-        except yeelight.BulbException as err:
+        except (asyncio.TimeoutError, yeelight.BulbException) as err:
             _LOGGER.error("Failed to get properties from %s: %s", host, err)
             raise CannotConnect from err
         _LOGGER.debug("Get properties: %s", bulb.last_properties)

tests/components/yeelight/test_config_flow.py

@@ -737,3 +737,47 @@ async def test_discovered_zeroconf(hass):
 
     assert result["type"] == RESULT_TYPE_ABORT
     assert result["reason"] == "already_configured"
+
+
+async def test_discovery_updates_ip(hass: HomeAssistant):
+    """Test discovery updtes ip."""
+    config_entry = MockConfigEntry(
+        domain=DOMAIN, data={CONF_HOST: "1.2.2.3"}, unique_id=ID
+    )
+    config_entry.add_to_hass(hass)
+
+    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_ZEROCONF},
+            data=ZEROCONF_DATA,
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == RESULT_TYPE_ABORT
+    assert result["reason"] == "already_configured"
+    assert config_entry.data[CONF_HOST] == IP_ADDRESS
+
+
+async def test_discovery_adds_missing_ip_id_only(hass: HomeAssistant):
+    """Test discovery adds missing ip."""
+    config_entry = MockConfigEntry(domain=DOMAIN, data={CONF_ID: ID})
+    config_entry.add_to_hass(hass)
+
+    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_ZEROCONF},
+            data=ZEROCONF_DATA,
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == RESULT_TYPE_ABORT
+    assert result["reason"] == "already_configured"
+    assert config_entry.data[CONF_HOST] == IP_ADDRESS