Bug fix entity creation (#27)

* Entity migration
arjenbos · Nov 23, 2024 · 102927d · 102927d
1 parent 55d13e9
commit 102927d
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 16 deletions.
diff --git a/custom_components/postnl/__init__.py b/custom_components/postnl/__init__.py
@@ -17,6 +17,7 @@
 from .const import DOMAIN, PLATFORMS
 from .graphql import PostNLGraphql
 from .login_api import PostNLLoginAPI
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -54,6 +55,30 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> True:
 
     hass.data[DOMAIN][entry.entry_id]['userinfo'] = userinfo
 
+    device_registry = dr.async_get(hass)
+    entity_registry = er.async_get(hass)
+
+    for device_entry in dr.async_entries_for_config_entry(
+        device_registry, entry.entry_id
+    ):
+        if (
+            device_entry.identifiers == {(DOMAIN, userinfo.get('account_id'))}
+        ):
+            _LOGGER.debug(
+                "Migrating entry %s"
+            )
+            for entity_entry in er.async_entries_for_device(
+                entity_registry, device_entry.id, True
+            ):
+                if entity_entry.unique_id.startswith(userinfo.get('account_id')):
+                    continue
+
+                unique_id_parts = entity_entry.unique_id.split("_")
+                entity_new_unique_id = userinfo.get('account_id') + "_" + unique_id_parts[1]
+                entity_registry.async_update_entity(
+                    entity_id=entity_entry.entity_id, new_unique_id=entity_new_unique_id
+                )
+
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     return True

diff --git a/custom_components/postnl/sensor.py b/custom_components/postnl/sensor.py
@@ -23,32 +23,18 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry, async_add_e
     await coordinator.async_config_entry_first_refresh()
     userinfo = hass.data[DOMAIN][entry.entry_id]['userinfo']
 
-    # Backwards compatibility support with old sensor names and to support the legacy lovelace.
-    unique_id_prefix = "postnl_"
-    entity_registry = async_get_entity_registry(hass)
-    if entity_registry.async_get_entity_id(
-        domain="sensor",
-        platform="postnl",
-        unique_id="postnl_delivery"
-    ) or entity_registry.async_get_entity_id(
-        domain="sensor",
-        platform="postnl",
-        unique_id="postnl_distribution"
-    ):
-        unique_id_prefix = userinfo.get('account_id') + "_"
-
     async_add_entities([
         PostNLDelivery(
             coordinator=coordinator,
             postnl_userinfo=userinfo,
-            unique_id= unique_id_prefix + "delivery",
+            unique_id= userinfo.get('account_id') + "_" + "delivery",
             name="PostNL_delivery"
         ),
         PostNLDelivery(
             coordinator=coordinator,
             postnl_userinfo=userinfo,
             name="PostNL_distribution",
-            unique_id=unique_id_prefix + "distribution",
+            unique_id=userinfo.get('account_id') + "_" + "distribution",
             receiver=False
         )
     ])