Merge pull request #671 from plugwise/rework-runtime-update-2

bouwew · web-flow · commit 1a62d80694d8 · 2024-06-22T11:10:34.000+02:00
Refactor runtime-update/remove related code
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,8 +2,9 @@
 
 ## Versions from 0.40 and up
 
-## Ongoing
+## v0.51.4
 
+- Refactor runtime-update/remove related code.
 - Implement various new Core features.
 
 ## v0.51.3
diff --git a/custom_components/plugwise/binary_sensor.py b/custom_components/plugwise/binary_sensor.py
@@ -98,33 +98,27 @@ async def async_setup_entry(
 
     @callback
     def _add_entities() -> None:
-        """Add Entities."""
+        """Add Entities during init and runtime."""
         if not coordinator.new_devices:
             return
 
         entities: list[PlugwiseBinarySensorEntity] = []
-        for device_id, device in coordinator.data.devices.items():
+        for device_id in coordinator.new_devices:
+            device = coordinator.data.devices[device_id]
             if not (binary_sensors := device.get(BINARY_SENSORS)):
                 continue
             for description in PLUGWISE_BINARY_SENSORS:
                 if description.key not in binary_sensors:
                     continue
-                entities.append(
-                    PlugwiseBinarySensorEntity(
-                        coordinator,
-                        device_id,
-                        description,
-                    )
-                )
+                entities.append(PlugwiseBinarySensorEntity(coordinator, device_id, description))
                 LOGGER.debug(
                     "Add %s %s binary sensor", device[ATTR_NAME], description.translation_key
                 )
 
         async_add_entities(entities)
 
-    entry.async_on_unload(coordinator.async_add_listener(_add_entities))
-
     _add_entities()
+    entry.async_on_unload(coordinator.async_add_listener(_add_entities))
 
 
 class PlugwiseBinarySensorEntity(PlugwiseEntity, BinarySensorEntity):
diff --git a/custom_components/plugwise/button.py b/custom_components/plugwise/button.py
@@ -36,31 +36,25 @@ async def async_setup_entry(
 
     @callback
     def _add_entities() -> None:
-        """Add Entities."""
+        """Add Entities during init and runtime."""
         if not coordinator.new_devices:
             return
 
         entities: list[PlugwiseButtonEntity] = []
         gateway = coordinator.data.gateway
-        for device_id, device in coordinator.data.devices.items():
+        for device_id in coordinator.new_devices:
+            device = coordinator.data.devices[device_id]
             if device_id == gateway[GATEWAY_ID] and REBOOT in gateway:
                 for description in BUTTON_TYPES:
-                    entities.append(
-                        PlugwiseButtonEntity(
-                            coordinator,
-                            device_id,
-                            description,
-                        )
-                    )
+                    entities.append(PlugwiseButtonEntity(coordinator, device_id, description))
                     LOGGER.debug(
                         "Add %s %s button", device[ATTR_NAME], description.key
                 )
 
         async_add_entities(entities)
 
-    entry.async_on_unload(coordinator.async_add_listener(_add_entities))
-
     _add_entities()
+    entry.async_on_unload(coordinator.async_add_listener(_add_entities))
 
 
 class PlugwiseButtonEntity(PlugwiseEntity, ButtonEntity):
diff --git a/custom_components/plugwise/climate.py b/custom_components/plugwise/climate.py
@@ -14,7 +14,13 @@
     HVACAction,
     HVACMode,
 )
-from homeassistant.const import ATTR_TEMPERATURE, STATE_OFF, STATE_ON, UnitOfTemperature
+from homeassistant.const import (
+    ATTR_NAME,
+    ATTR_TEMPERATURE,
+    STATE_OFF,
+    STATE_ON,
+    UnitOfTemperature,
+)
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
@@ -33,6 +39,7 @@
     GATEWAY_ID,
     HEATING_STATE,
     LOCATION,
+    LOGGER,
     LOWER_BOUND,
     MASTER_THERMOSTATS,
     MODE,
@@ -57,29 +64,33 @@ async def async_setup_entry(
     entry: PlugwiseConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
-    """Set up the Smile Thermostats from a ConfigEntry."""
+    """Set up the Plugwise thermostats from a ConfigEntry."""
     coordinator = entry.runtime_data
     homekit_enabled: bool = entry.options.get(
         CONF_HOMEKIT_EMULATION, False
     )  # pw-beta homekit emulation
 
     @callback
     def _add_entities() -> None:
-        """Add Entities."""
+        """Add Entities during init and runtime."""
         if not coordinator.new_devices:
             return
 
-        async_add_entities(
-            PlugwiseClimateEntity(
-                coordinator, device_id, homekit_enabled
-            )  # pw-beta homekit emulation
-            for device_id, device in coordinator.data.devices.items()
-            if device[DEV_CLASS] in MASTER_THERMOSTATS
-        )
+        entities: list[PlugwiseClimateEntity] = []
+        for device_id in coordinator.new_devices:
+            device = coordinator.data.devices[device_id]
+            if device[DEV_CLASS] in MASTER_THERMOSTATS:
+                entities.append(
+                    PlugwiseClimateEntity(
+                        coordinator, device_id, homekit_enabled
+                    )  # pw-beta homekit emulation
+                )
+                LOGGER.debug("Add climate %s", device[ATTR_NAME])
 
-    entry.async_on_unload(coordinator.async_add_listener(_add_entities))
+        async_add_entities(entities)
 
     _add_entities()
+    entry.async_on_unload(coordinator.async_add_listener(_add_entities))
 
 
 class PlugwiseClimateEntity(PlugwiseEntity, ClimateEntity):
diff --git a/custom_components/plugwise/coordinator.py b/custom_components/plugwise/coordinator.py
@@ -24,7 +24,6 @@
 from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.debounce import Debouncer
-from homeassistant.helpers.device_registry import DeviceEntry, DeviceRegistry
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
 from .const import (
@@ -37,45 +36,6 @@
 )
 
 
-async def cleanup_device_and_entity_registry(
-    data: PlugwiseData,
-    device_reg: DeviceRegistry,
-    device_list: list[DeviceEntry],
-    entry: ConfigEntry,
-) -> None:
-    """Remove deleted devices from device- and entity-registry."""
-    if len(device_list) - len(data.devices.keys()) <= 0:
-        return
-
-    # via_device cannot be None, this will result in the deletion
-    # of other Plugwise Gateways when present!
-    via_device: str = ""
-    for device_entry in device_list:
-        if not device_entry.identifiers:
-            continue  # pragma: no cover
-
-        item = list(list(device_entry.identifiers)[0])
-        if item[0] != DOMAIN:
-            continue  # pragma: no cover
-
-        # First find the Plugwise via_device, this is always the first device
-        if item[1] == data.gateway[GATEWAY_ID]:
-            via_device = device_entry.id
-        elif ( # then remove the connected orphaned device(s)
-            device_entry.via_device_id == via_device
-            and item[1] not in list(data.devices.keys())
-        ):
-            device_reg.async_update_device(
-                device_entry.id, remove_config_entry_id=entry.entry_id
-            )
-            LOGGER.debug(
-                "Removed %s device %s %s from device_registry",
-                DOMAIN,
-                device_entry.model,
-                item[1],
-            )
-
-
 class PlugwiseDataUpdateCoordinator(DataUpdateCoordinator[PlugwiseData]):
     """Class to manage fetching Plugwise data from single endpoint."""
 
@@ -114,9 +74,8 @@ def __init__(
             timeout=30,
             websession=async_get_clientsession(hass, verify_ssl=False),
         )
-        self._unavailable_logged = False
-        self.device_list: list[DeviceEntry] = []
-        self.new_devices: bool = False
+        self._current_devices: set[str] = set()
+        self.new_devices: set[str] = set()
         self.update_interval = update_interval
 
     async def _connect(self) -> None:
@@ -155,17 +114,46 @@ async def _async_update_data(self) -> PlugwiseData:
                 raise ConfigEntryError("Device with unsupported firmware") from err
         else:
             LOGGER.debug(f"{self.api.smile_name} data: %s", data)
-            device_reg = dr.async_get(self.hass)
-            device_list = dr.async_entries_for_config_entry(
-                device_reg, self.config_entry.entry_id
-            )
-            await cleanup_device_and_entity_registry(
-                data,
-                device_reg,
-                device_list,
-                self.config_entry
-            )
-            self.new_devices = len(data.devices.keys()) - len(self.device_list) > 0
-            self.device_list = device_list
+            self._async_add_remove_devices(data, self.config_entry)
 
         return data
+
+    def _async_add_remove_devices(self, data: PlugwiseData, entry: ConfigEntry) -> None:
+        """Add new Plugwise devices, remove non-existing devices."""
+        # Check for new or removed devices
+        self.new_devices = set(data.devices) - self._current_devices
+        removed_devices = self._current_devices - set(data.devices)
+        self._current_devices = set(data.devices)
+
+        if removed_devices:
+            self._async_remove_devices(data, entry)
+
+    def _async_remove_devices(self, data: PlugwiseData, entry: ConfigEntry) -> 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
+        )
+        # via_device cannot be None, this will result in the deletion
+        # of other Plugwise Gateways when present!
+        via_device: str = ""
+        for device_entry in device_list:
+            if device_entry.identifiers:
+                item = list(list(device_entry.identifiers)[0])
+                if item[0] == DOMAIN:
+                    # First find the Plugwise via_device, this is always the first device
+                    if item[1] == data.gateway[GATEWAY_ID]:
+                        via_device = device_entry.id
+                    elif ( # then remove the connected orphaned device(s)
+                        device_entry.via_device_id == via_device
+                        and item[1] not in list(data.devices.keys())
+                    ):
+                        device_reg.async_update_device(
+                            device_entry.id, remove_config_entry_id=entry.entry_id
+                        )
+                        LOGGER.debug(
+                            "Removed %s device %s %s from device_registry",
+                            DOMAIN,
+                            device_entry.model,
+                            item[1],
+                        )
diff --git a/custom_components/plugwise/manifest.json b/custom_components/plugwise/manifest.json
@@ -9,5 +9,5 @@
   "iot_class": "local_polling",
   "loggers": ["plugwise"],
   "requirements": ["plugwise==0.38.3"],
-  "version": "0.51.3"
+  "version": "0.51.4"
 }
diff --git a/custom_components/plugwise/number.py b/custom_components/plugwise/number.py
@@ -71,12 +71,13 @@ async def async_setup_entry(
 
     @callback
     def _add_entities() -> None:
-        """Add Entities."""
+        """Add Entities during init and runtime."""
         if not coordinator.new_devices:
             return
 
         entities: list[PlugwiseNumberEntity] = []
-        for device_id, device in coordinator.data.devices.items():
+        for device_id in coordinator.new_devices:
+            device = coordinator.data.devices[device_id]
             for description in NUMBER_TYPES:
                 if description.key in device:
                     entities.append(
@@ -88,9 +89,8 @@ def _add_entities() -> None:
 
         async_add_entities(entities)
 
-    entry.async_on_unload(coordinator.async_add_listener(_add_entities))
-
     _add_entities()
+    entry.async_on_unload(coordinator.async_add_listener(_add_entities))
 
 
 class PlugwiseNumberEntity(PlugwiseEntity, NumberEntity):
diff --git a/custom_components/plugwise/select.py b/custom_components/plugwise/select.py
@@ -73,17 +73,18 @@ async def async_setup_entry(
     entry: PlugwiseConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
-    """Set up the Smile selector from a ConfigEntry."""
+    """Set up the Plugwise selectors from a ConfigEntry."""
     coordinator = entry.runtime_data
 
     @callback
     def _add_entities() -> None:
-        """Add Entities."""
+        """Add Entities during init and runtime."""
         if not coordinator.new_devices:
             return
 
         entities: list[PlugwiseSelectEntity] = []
-        for device_id, device in coordinator.data.devices.items():
+        for device_id in coordinator.new_devices:
+            device = coordinator.data.devices[device_id]
             for description in SELECT_TYPES:
                 if description.options_key in device:
                     entities.append(
@@ -95,9 +96,8 @@ def _add_entities() -> None:
 
         async_add_entities(entities)
 
-    entry.async_on_unload(coordinator.async_add_listener(_add_entities))
-
     _add_entities()
+    entry.async_on_unload(coordinator.async_add_listener(_add_entities))
 
 
 class PlugwiseSelectEntity(PlugwiseEntity, SelectEntity):
diff --git a/custom_components/plugwise/sensor.py b/custom_components/plugwise/sensor.py
@@ -458,38 +458,32 @@ async def async_setup_entry(
     entry: PlugwiseConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
-    """Set up the Smile sensors from a ConfigEntry."""
+    """Set up the Plugwise sensors from a ConfigEntry."""
     coordinator = entry.runtime_data
 
     @callback
     def _add_entities() -> None:
-        """Add Entities."""
+        """Add Entities during init and runtime."""
         if not coordinator.new_devices:
             return
 
         entities: list[PlugwiseSensorEntity] = []
-        for device_id, device in coordinator.data.devices.items():
+        for device_id in coordinator.new_devices:
+            device = coordinator.data.devices[device_id]
             if not (sensors := device.get(SENSORS)):
                 continue
             for description in PLUGWISE_SENSORS:
                 if description.key not in sensors:
                     continue
-                entities.append(
-                    PlugwiseSensorEntity(
-                        coordinator,
-                        device_id,
-                        description,
-                    )
-                )
+                entities.append(PlugwiseSensorEntity(coordinator, device_id, description))
                 LOGGER.debug(
                     "Add %s %s sensor", device[ATTR_NAME], description.translation_key or description.key
                 )
 
         async_add_entities(entities)
 
-    entry.async_on_unload(coordinator.async_add_listener(_add_entities))
-
     _add_entities()
+    entry.async_on_unload(coordinator.async_add_listener(_add_entities))
 
 
 class PlugwiseSensorEntity(PlugwiseEntity, SensorEntity):
diff --git a/custom_components/plugwise/switch.py b/custom_components/plugwise/switch.py
diff --git a/tests/components/plugwise/test_init.py b/tests/components/plugwise/test_init.py

Original file line number	Diff line number	Diff line change
`@@ -9,5 +9,5 @@`
`9`	`9`	`"iot_class": "local_polling",`
`10`	`10`	`"loggers": ["plugwise"],`
`11`	`11`	`"requirements": ["plugwise==0.38.3"],`
`12`		`- "version": "0.51.3"`
	`12`	`+ "version": "0.51.4"`
`13`	`13`	`}`