Improve Plugwise DataUpdateCoordinator (#165715)

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

homeassistant/components/plugwise/coordinator.py

@@ -36,11 +36,6 @@ type PlugwiseConfigEntry = ConfigEntry[PlugwiseDataUpdateCoordinator]
 class PlugwiseDataUpdateCoordinator(DataUpdateCoordinator[dict[str, GwEntityData]]):
     """Class to manage fetching Plugwise data from single endpoint."""
 
-    _connected: bool = False
-    _current_devices: set[str]
-    _stored_devices: set[str]
-    new_devices: set[str]
-
     config_entry: PlugwiseConfigEntry
 
     def __init__(self, hass: HomeAssistant, config_entry: PlugwiseConfigEntry) -> None:
@@ -68,9 +63,10 @@ class PlugwiseDataUpdateCoordinator(DataUpdateCoordinator[dict[str, GwEntityData
             port=self.config_entry.data.get(CONF_PORT, DEFAULT_PORT),
             websession=async_get_clientsession(hass, verify_ssl=False),
         )
-        self._current_devices = set()
-        self._stored_devices = set()
-        self.new_devices = set()
+        self._connected: bool = False
+        self._current_devices: set[str] = set()
+        self._stored_devices: set[str] = set()
+        self.new_devices: set[str] = set()
 
     async def _connect(self) -> None:
         """Connect to the Plugwise Smile.
@@ -132,10 +128,10 @@ class PlugwiseDataUpdateCoordinator(DataUpdateCoordinator[dict[str, GwEntityData
                 translation_key="unsupported_firmware",
             ) from err
 
-        self._async_add_remove_devices(data)
+        self._add_remove_devices(data)
         return data
 
-    def _async_add_remove_devices(self, data: dict[str, GwEntityData]) -> None:
+    def _add_remove_devices(self, data: dict[str, GwEntityData]) -> None:
         """Add new Plugwise devices, remove non-existing devices."""
         set_of_data = set(data)
         # Check for new or removed devices,
@@ -146,35 +142,28 @@ class PlugwiseDataUpdateCoordinator(DataUpdateCoordinator[dict[str, GwEntityData
             self._stored_devices if not self._current_devices else self._current_devices
         )
         self._current_devices = set_of_data
-        if current_devices - set_of_data:  # device(s) to remove
-            self._async_remove_devices(data)
+        if removed_devices := (current_devices - set_of_data):  # device(s) to remove
+            self._remove_devices(removed_devices)
 
-    def _async_remove_devices(self, data: dict[str, GwEntityData]) -> None:
+    def _remove_devices(self, removed_devices: set[str]) -> None:
         """Clean registries when removed devices found."""
         device_reg = dr.async_get(self.hass)
-        device_list = dr.async_entries_for_config_entry(
-            device_reg, self.config_entry.entry_id
-        )
+        for device_id in removed_devices:
+            device_entry = device_reg.async_get_device({(DOMAIN, device_id)})
+            if device_entry is None:
+                LOGGER.warning(
+                    "Failed to remove %s device/zone %s, not present in device_registry",
+                    DOMAIN,
+                    device_id,
+                )
+                continue  # pragma: no cover
 
-        # First find the Plugwise via_device
-        gateway_device = device_reg.async_get_device({(DOMAIN, self.api.gateway_id)})
-        assert gateway_device is not None
-        via_device_id = gateway_device.id
-        # Then remove the connected orphaned device(s)
-        for device_entry in device_list:
-            for identifier in device_entry.identifiers:
-                if (
-                    identifier[0] == DOMAIN
-                    and device_entry.via_device_id == via_device_id
-                    and identifier[1] not in data
-                ):
-                    device_reg.async_update_device(
-                        device_entry.id,
-                        remove_config_entry_id=self.config_entry.entry_id,
-                    )
-                    LOGGER.debug(
-                        "Removed %s device/zone %s %s from device_registry",
-                        DOMAIN,
-                        device_entry.model,
-                        identifier[1],
-                    )
+            device_reg.async_update_device(
+                device_entry.id, remove_config_entry_id=self.config_entry.entry_id
+            )
+            LOGGER.debug(
+                "%s %s %s removed from device_registry",
+                DOMAIN,
+                device_entry.model,
+                device_id,
+            )

tests/components/plugwise/test_init.py

@@ -305,10 +305,10 @@ async def test_update_device(
             )
             == 12
         )
-        item_list: list[str] = []
-        for device_entry in list(device_registry.devices.values()):
-            item_list.extend(x[1] for x in device_entry.identifiers)
-        assert "01234567890abcdefghijklmnopqrstu" in item_list
+        device_entry = device_registry.async_get_device(
+            identifiers={(DOMAIN, "01234567890abcdefghijklmnopqrstu")}
+        )
+        assert device_entry is not None
 
     # Remove the existing Tom/Floor
     data["f871b8c4d63549319221e294e4f88074"]["thermostats"].update(
@@ -336,10 +336,10 @@ async def test_update_device(
             )
             == 11
         )
-        item_list: list[str] = []
-        for device_entry in list(device_registry.devices.values()):
-            item_list.extend(x[1] for x in device_entry.identifiers)
-        assert "1772a4ea304041adb83f357b751341ff" not in item_list
+        device_entry = device_registry.async_get_device(
+            identifiers={(DOMAIN, "1772a4ea304041adb83f357b751341ff")}
+        )
+        assert device_entry is None
 
 
 @pytest.mark.parametrize("chosen_env", ["m_adam_heating"], indirect=True)
@@ -350,24 +350,26 @@ async def test_delete_removed_device(
     mock_smile_adam_heat_cool: MagicMock,
     device_registry: dr.DeviceRegistry,
     init_integration: MockConfigEntry,
+    freezer: FrozenDateTimeFactory,
 ) -> None:
-    """Test device removal at integration init."""
+    """Test device removal after a coordinator refresh."""
     data = mock_smile_adam_heat_cool.async_update.return_value
 
-    item_list: list[str] = []
-    for device_entry in device_registry.devices.values():
-        item_list.extend(x[1] for x in device_entry.identifiers)
-    assert "14df5c4dc8cb4ba69f9d1ac0eaf7c5c6" in item_list
+    device_entry = device_registry.async_get_device(
+        identifiers={(DOMAIN, "14df5c4dc8cb4ba69f9d1ac0eaf7c5c6")}
+    )
+    assert device_entry is not None
 
     data.pop("14df5c4dc8cb4ba69f9d1ac0eaf7c5c6")
     with patch(HA_PLUGWISE_SMILE_ASYNC_UPDATE, return_value=data):
-        await hass.config_entries.async_reload(init_integration.entry_id)
+        freezer.tick(timedelta(minutes=1))
+        async_fire_time_changed(hass)
         await hass.async_block_till_done()
 
-    item_list = []
-    for device_entry in device_registry.devices.values():
-        item_list.extend(x[1] for x in device_entry.identifiers)
-    assert "14df5c4dc8cb4ba69f9d1ac0eaf7c5c6" not in item_list
+    device_entry = device_registry.async_get_device(
+        identifiers={(DOMAIN, "14df5c4dc8cb4ba69f9d1ac0eaf7c5c6")}
+    )
+    assert device_entry is None
 
 
 @pytest.mark.parametrize("chosen_env", ["m_adam_heating"], indirect=True)