Final code improvement, also improve test-case code

plugwise · Mar 8, 2024 · e7606f2 · e7606f2
1 parent f50987b
commit e7606f2
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 8 deletions.
diff --git a/custom_components/plugwise/__init__.py b/custom_components/plugwise/__init__.py
@@ -94,27 +94,28 @@ def cleanup_device_registry(
     data: PlugwiseData,
 ) -> None:
     """Remove deleted devices from device-registry."""
-    gateway_id = data.gateway["gateway_id"]
-    plugwise_device_list = list(data.devices.keys()).remove(gateway_id)
+    gateway_id: str = data.gateway["gateway_id"]
+    plugwise_device_list: list[str] = list(data.devices.keys())
     device_registry = dr.async_get(hass)
     dev_reg_list = list(device_registry.devices.items())
-    # First find the Plugwise via_device_id   
+    # First find the Plugwise via_device_id
     for dev_id, device_entry in dev_reg_list:
         item = list(list(device_entry.identifiers)[0])
         if item[0] == DOMAIN and item[1] == gateway_id:
             via_device = dev_id
 
     # Find and remove the orphaned device(s) connected to the via_device
-    for _, device_entry in dev_reg_list:
+    for dev_id, device_entry in dev_reg_list:
         item = list(list(device_entry.identifiers)[0])
         if (
             item[0] == DOMAIN
-            and item[1] not in plugwise_device_list
             and device_entry.via_device_id == via_device
+            and item[1] not in plugwise_device_list
+
         ):
-            # device_registry.async_remove_device(dev_id)
+            device_registry.async_remove_device(dev_id)
             LOGGER.debug(
-                "HOI found orphaned device: %s %s %s in device_registry",
+                "Removed %s device %s %s from device_registry",
                 DOMAIN,
                 device_entry.model,
                 item[1],

diff --git a/tests/components/plugwise/test_init.py b/tests/components/plugwise/test_init.py
@@ -216,7 +216,7 @@ async def test_device_removal(
         await hass.config_entries.async_reload(mock_config_entry.entry_id)
         await hass.async_block_till_done()
 
-    devices = dr.async_entries_for_config_entry(dev_reg, mock_config_entry.entry_id)
+    dev_reg = dr.async_get(hass)
     item_list = []
     for device_entry in list(dev_reg.devices.values()):
         for item in device_entry.identifiers: