diff --git a/homeassistant/components/yeelight/__init__.py b/homeassistant/components/yeelight/__init__.py index 19fe5c550ad..a4ff947191e 100644 --- a/homeassistant/components/yeelight/__init__.py +++ b/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 diff --git a/tests/components/yeelight/test_init.py b/tests/components/yeelight/test_init.py index 4b3ac8e0e83..cee798308c4 100644 --- a/tests/components/yeelight/test_init.py +++ b/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(