diff --git a/custom_components/localtuya/__init__.py b/custom_components/localtuya/__init__.py
index 78a9443..b126c2f 100644
--- a/custom_components/localtuya/__init__.py
+++ b/custom_components/localtuya/__init__.py
@@ -139,15 +139,16 @@ async def async_setup(hass: HomeAssistant, config: dict):
             )
             new_data[ATTR_UPDATED_AT] = str(int(time.time() * 1000))
             hass.config_entries.async_update_entry(entry, data=new_data)
-            device = hass.data[DOMAIN][TUYA_DEVICES][device_id]
-            if not device.connected:
-                device.async_connect()
+   
         elif device_id in hass.data[DOMAIN][TUYA_DEVICES]:
-            # _LOGGER.debug("Device %s found with IP %s", device_id, device_ip)
+            _LOGGER.debug("Device %s found with IP %s", device_id, device_ip)
+
+        device = hass.data[DOMAIN][TUYA_DEVICES].get(device_id)
+        if not device:
+            _LOGGER.warning(f"Could not find device for device_id {device_id}")
+        elif not device.connected:
+            device.async_connect()
 
-            device = hass.data[DOMAIN][TUYA_DEVICES][device_id]
-            if not device.connected:
-                device.async_connect()
 
     def _shutdown(event):
         """Clean up resources when shutting down."""
diff --git a/custom_components/localtuya/common.py b/custom_components/localtuya/common.py
index cd503c2..4ad46f6 100644
--- a/custom_components/localtuya/common.py
+++ b/custom_components/localtuya/common.py
@@ -124,8 +124,10 @@ def async_config_entry_by_device_id(hass, device_id):
     """Look up config entry by device id."""
     current_entries = hass.config_entries.async_entries(DOMAIN)
     for entry in current_entries:
-        if device_id in entry.data[CONF_DEVICES]:
+        if device_id in entry.data.get(CONF_DEVICES, []):
             return entry
+        else:
+            _LOGGER.warning(f"Missing device configuration for device_id {device_id}")
     return None