Fix unload KeyError for failed device setup (#5381)

* fix: check device existence in async_delete_device before unloading

* docs: explain device cache cleanup guard

custom_components/tuya_local/__init__.py

@@ -1003,9 +1003,15 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry):
 
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry):
-    _LOGGER.debug("Unloading entry for device: %s", get_device_id(entry.data))
+    device_id = get_device_id(entry.data)
+    _LOGGER.debug("Unloading entry for device: %s", device_id)
     config = entry.data
-    data = hass.data[DOMAIN][get_device_id(config)]
+    domain_data = hass.data.get(DOMAIN, {})
+    data = domain_data.get(device_id)
+    if data is None:
+        await async_delete_device(hass, config)
+        return True
+
     device_conf = await hass.async_add_executor_job(
         get_config,
         config[CONF_TYPE],
@@ -1023,7 +1029,7 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry):
         await hass.config_entries.async_forward_entry_unload(entry, e)
 
     await async_delete_device(hass, config)
-    del hass.data[DOMAIN][get_device_id(config)]
+    domain_data.pop(device_id, None)
 
     return True
 

custom_components/tuya_local/device.py

@@ -813,7 +813,19 @@ def setup_device(hass: HomeAssistant, config: dict):
 async def async_delete_device(hass: HomeAssistant, config: dict):
     device_id = get_device_id(config)
     _LOGGER.info("Deleting device: %s", device_id)
-    await hass.data[DOMAIN][device_id]["device"].async_stop()
-    del hass.data[DOMAIN][device_id]["device"]
-    del hass.data[DOMAIN][device_id]["tuyadevice"]
-    del hass.data[DOMAIN][device_id]["tuyadevicelock"]
+    domain_data = hass.data.get(DOMAIN, {})
+    device_entry = domain_data.get(device_id)
+    if device_entry is None:
+        return
+
+    device = device_entry.get("device")
+    if device is not None:
+        await device.async_stop()
+        device_entry.pop("device", None)
+    device_entry.pop("tuyadevice", None)
+    device_entry.pop("tuyadevicelock", None)
+    # Platform setup may cache entity instances in this bucket by config_id.
+    # Only drop empty buckets here; async_unload_entry removes the whole bucket
+    # after forwarded platform unloads complete.
+    if not device_entry:
+        domain_data.pop(device_id, None)

tests/test_config_flow.py

@@ -12,6 +12,7 @@ from pytest_homeassistant_custom_component.common import MockConfigEntry
 from custom_components.tuya_local import (
     async_migrate_entry,
     async_setup_entry,
+    async_unload_entry,
     config_flow,
     get_device_unique_id,
 )
@@ -127,6 +128,30 @@ async def test_async_setup_entry_cleans_up_failed_device(hass, mocker, refresh_e
     mock_device._api.set_socketPersistent.assert_called_with(False)
 
 
+@pytest.mark.asyncio
+async def test_async_unload_entry_ignores_missing_device_data(hass):
+    """Unload should tolerate entries that failed before device data was cached."""
+
+    hass.data[DOMAIN] = {}
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        version=13,
+        minor_version=18,
+        title="test",
+        data={
+            CONF_DEVICE_ID: "deviceid",
+            CONF_HOST: "hostname",
+            CONF_LOCAL_KEY: TESTKEY,
+            CONF_POLL_ONLY: False,
+            CONF_PROTOCOL_VERSION: 3.4,
+            CONF_TYPE: "kogan_kahtp_heater",
+        },
+        options={},
+    )
+
+    assert await async_unload_entry(hass, entry)
+
+
 @pytest.mark.asyncio
 async def test_migrate_entry(hass, mocker):
     """Test migration from old entry format."""

tests/test_device.py

@@ -4,7 +4,8 @@ from time import time
 import pytest
 
 # from homeassistant.const import EVENT_HOMEASSISTANT_STARTED, EVENT_HOMEASSISTANT_STOP
-from custom_components.tuya_local.device import TuyaLocalDevice
+from custom_components.tuya_local.const import CONF_DEVICE_ID, DOMAIN
+from custom_components.tuya_local.device import TuyaLocalDevice, async_delete_device
 
 from .const import EUROM_600_HEATER_PAYLOAD
 
@@ -68,6 +69,25 @@ def test_device_info(subject, mock_api):
     }
 
 
+@pytest.mark.asyncio
+async def test_delete_keeps_device_entry_when_stop_fails(hass, mocker):
+    """Device cache should not be removed before stop succeeds."""
+    device = mocker.MagicMock()
+    device.async_stop = mocker.AsyncMock(side_effect=RuntimeError("stop failed"))
+    hass.data[DOMAIN] = {
+        "deviceid": {
+            "device": device,
+            "tuyadevice": mocker.MagicMock(),
+            "tuyadevicelock": mocker.MagicMock(),
+        }
+    }
+
+    with pytest.raises(RuntimeError, match="stop failed"):
+        await async_delete_device(hass, {CONF_DEVICE_ID: "deviceid"})
+
+    assert hass.data[DOMAIN]["deviceid"]["device"] is device
+
+
 def test_has_returned_state(subject):
     """Returns True if the device has returned its state."""
     subject._cached_state = EUROM_600_HEATER_PAYLOAD