Ensure yeelight can be unloaded when device is offline (#56464)

homeassistant/components/yeelight/__init__.py

@@ -303,18 +303,21 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
-    data_config_entries = hass.data[DOMAIN][DATA_CONFIG_ENTRIES]
-    entry_data = data_config_entries[entry.entry_id]
-
-    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 = hass.data[DOMAIN][DATA_CONFIG_ENTRIES]
+    if entry.entry_id not in data_config_entries:
+        # Device not online
+        return True
+
+    entry_data = data_config_entries[entry.entry_id]
+    unload_ok = True
+    if entry_data[DATA_PLATFORMS_LOADED]:
+        unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
+
     if DATA_DEVICE in entry_data:
         device = entry_data[DATA_DEVICE]
         _LOGGER.debug("Shutting down Yeelight Listener")
@@ -322,8 +325,7 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         _LOGGER.debug("Yeelight Listener stopped")
 
     data_config_entries.pop(entry.entry_id)
-
-    return True
+    return unload_ok
 
 
 @callback

tests/components/yeelight/test_init.py

@@ -366,6 +366,26 @@ async def test_async_listen_error_late_discovery(hass, caplog):
     assert config_entry.options[CONF_MODEL] == MODEL
 
 
+async def test_unload_before_discovery(hass, caplog):
+    """Test unloading before discovery."""
+    config_entry = MockConfigEntry(domain=DOMAIN, data=CONFIG_ENTRY_DATA)
+    config_entry.add_to_hass(hass)
+
+    mocked_bulb = _mocked_bulb(cannot_connect=True)
+
+    with _patch_discovery(no_device=True), 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.LOADED
+    await hass.config_entries.async_unload(config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert config_entry.state is ConfigEntryState.NOT_LOADED
+
+
 async def test_async_listen_error_has_host_with_id(hass: HomeAssistant):
     """Test the async listen error."""
     config_entry = MockConfigEntry(