Move ISY994 data to dataclass and remove bad entities (#85744)

homeassistant/components/isy994/util.py

@@ -1,69 +1,34 @@
 """ISY utils."""
 from __future__ import annotations
 
-from pyisy.constants import PROP_COMMS_ERROR, PROTO_INSTEON
+from homeassistant.core import HomeAssistant, callback
+import homeassistant.helpers.entity_registry as er
 
-from homeassistant.const import Platform
-from homeassistant.core import HomeAssistant
-
-from .const import (
-    CONF_NETWORK,
-    DOMAIN,
-    ISY_NET_RES,
-    ISY_NODES,
-    ISY_PROGRAMS,
-    ISY_ROOT,
-    ISY_ROOT_NODES,
-    ISY_VARIABLES,
-    NODE_PLATFORMS,
-    PROGRAM_PLATFORMS,
-    ROOT_NODE_PLATFORMS,
-    SENSOR_AUX,
-)
+from .const import _LOGGER, DOMAIN
 
 
-def unique_ids_for_config_entry_id(
-    hass: HomeAssistant, config_entry_id: str
-) -> set[tuple[Platform | str, str]]:
-    """Find all the unique ids for a config entry id."""
-    hass_isy_data = hass.data[DOMAIN][config_entry_id]
-    isy = hass_isy_data[ISY_ROOT]
-    current_unique_ids: set[tuple[Platform | str, str]] = {
-        (Platform.BUTTON, f"{isy.uuid}_query")
+@callback
+def _async_cleanup_registry_entries(hass: HomeAssistant, entry_id: str) -> None:
+    """Remove extra entities that are no longer part of the integration."""
+    entity_registry = er.async_get(hass)
+    isy_data = hass.data[DOMAIN][entry_id]
+
+    existing_entries = er.async_entries_for_config_entry(entity_registry, entry_id)
+    entities = {
+        (entity.domain, entity.unique_id): entity.entity_id
+        for entity in existing_entries
     }
 
-    # Structure and prefixes here must match what's added in __init__ and helpers
-    for platform in NODE_PLATFORMS:
-        for node in hass_isy_data[ISY_NODES][platform]:
-            current_unique_ids.add((platform, f"{isy.uuid}_{node.address}"))
+    extra_entities = set(entities.keys()).difference(isy_data.unique_ids)
+    if not extra_entities:
+        return
 
-    for node, control in hass_isy_data[ISY_NODES][SENSOR_AUX]:
-        current_unique_ids.add(
-            (Platform.SENSOR, f"{isy.uuid}_{node.address}_{control}")
-        )
-        current_unique_ids.add(
-            (Platform.SENSOR, f"{isy.uuid}_{node.address}_{PROP_COMMS_ERROR}")
-        )
+    for entity in extra_entities:
+        if entity_registry.async_is_registered(entities[entity]):
+            entity_registry.async_remove(entities[entity])
 
-    for platform in PROGRAM_PLATFORMS:
-        for _, node, _ in hass_isy_data[ISY_PROGRAMS][platform]:
-            current_unique_ids.add((platform, f"{isy.uuid}_{node.address}"))
-
-    for node, _ in hass_isy_data[ISY_VARIABLES][Platform.NUMBER]:
-        current_unique_ids.add((Platform.NUMBER, f"{isy.uuid}_{node.address}"))
-        current_unique_ids.add((Platform.NUMBER, f"{isy.uuid}_{node.address}_init"))
-    for _, node in hass_isy_data[ISY_VARIABLES][Platform.SENSOR]:
-        current_unique_ids.add((Platform.SENSOR, f"{isy.uuid}_{node.address}"))
-
-    for platform in ROOT_NODE_PLATFORMS:
-        for node in hass_isy_data[ISY_ROOT_NODES][platform]:
-            current_unique_ids.add((platform, f"{isy.uuid}_{node.address}_query"))
-            if platform == Platform.BUTTON and node.protocol == PROTO_INSTEON:
-                current_unique_ids.add((platform, f"{isy.uuid}_{node.address}_beep"))
-
-    for node in hass_isy_data[ISY_NET_RES]:
-        current_unique_ids.add(
-            (Platform.BUTTON, f"{isy.uuid}_{CONF_NETWORK}_{node.address}")
-        )
-
-    return current_unique_ids
+    _LOGGER.debug(
+        ("Cleaning up ISY entities: removed %s extra entities for config entry %s"),
+        len(extra_entities),
+        entry_id,
+    )