From 6faa47c670e511b15502250ce7ff768c7d86dce2 Mon Sep 17 00:00:00 2001
From: Bram <bgerritsen@emico.nl>
Date: Thu, 26 Dec 2024 10:50:18 +0100
Subject: [PATCH] fix: fully implement entity registry updating and 100%
 coverage

---
 .../sensors/group/tracked_untracked.py        | 19 +++++++++++------
 tests/sensors/group/test_tracked_untracked.py | 21 +++++++++++++------
 2 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/custom_components/powercalc/sensors/group/tracked_untracked.py b/custom_components/powercalc/sensors/group/tracked_untracked.py
index 16d571335..b1f0615ce 100644
--- a/custom_components/powercalc/sensors/group/tracked_untracked.py
+++ b/custom_components/powercalc/sensors/group/tracked_untracked.py
@@ -126,16 +126,23 @@ async def _handle_entity_registry_updated(
         entity_id = event.data["entity_id"]
         action = event.data["action"]
 
-        if action == "update":
-            return
-        if action == "remove" and entity_id not in self.tracked_entities:
-            return
+        if action == "update" and "old_entity_id" in event.data:
+            if event.data["old_entity_id"] in self.tracked_entities:
+                return await self.reload()
+            return None
+
+        if action == "remove" and entity_id in self.tracked_entities:
+            return await self.reload()
+
         if action == "create":
             registry = er.async_get(self.hass)
             entity_entry = registry.async_get(entity_id)
-            if entity_entry and entity_entry.original_device_class != SensorDeviceClass.POWER:
-                return
+            if entity_entry and entity_entry.original_device_class == SensorDeviceClass.POWER:
+                return await self.reload()
+        return None
 
+    async def reload(self) -> None:
+        """Reload the config entry."""
         await self.hass.config_entries.async_reload(self.config_entry.entry_id)
 
     async def create_tracked_power_sensor(
diff --git a/tests/sensors/group/test_tracked_untracked.py b/tests/sensors/group/test_tracked_untracked.py
index c48e3320f..53a3cfc1d 100644
--- a/tests/sensors/group/test_tracked_untracked.py
+++ b/tests/sensors/group/test_tracked_untracked.py
@@ -130,7 +130,11 @@ async def test_auto_tracking_entities(hass: HomeAssistant) -> None:
 
 async def test_entity_registry_updates(hass: HomeAssistant) -> None:
     """Test that the tracked power sensor is updated when power sensors are added or removed to the system"""
-    mock_sensors_in_registry(hass, ["sensor.test1_power", "sensor.test2_power"])
+    mock_sensors_in_registry(
+        hass,
+        ["sensor.test1_power", "sensor.test2_power", "sensor.test3_power"],
+        ["sensor.test1_energy"],
+    )
     entity_registry = er.async_get(hass)
     create_mock_entry(
         hass,
@@ -146,17 +150,22 @@ async def test_entity_registry_updates(hass: HomeAssistant) -> None:
 
     hass.states.async_set("sensor.test1_power", "10")
     hass.states.async_set("sensor.test2_power", "5")
+    hass.states.async_set("sensor.test3_power", "10")
     await hass.async_block_till_done()
 
     tracked_power_sensor = hass.data[DOMAIN][DATA_GROUP_ENTITIES]["sensor.tracked_power"]
-    assert tracked_power_sensor.entities == {"sensor.test1_power", "sensor.test2_power"}
-    assert hass.states.get("sensor.tracked_power").state == "15.00"
+    assert tracked_power_sensor.entities == {"sensor.test1_power", "sensor.test2_power", "sensor.test3_power"}
+    assert hass.states.get("sensor.tracked_power").state == "25.00"
 
     # Remove one of the tracked entities from registry
     entity_registry.async_remove("sensor.test2_power")
+    entity_registry.async_remove("sensor.test1_energy")  # for coverage
+    # Change the entity_id of one of the tracked entities
+    entity_registry.async_update_entity("sensor.test1_power", new_entity_id="sensor.test1_power_new")
+    entity_registry.async_update_entity("sensor.test3_power", icon="mdi:power")  # irrelevant change for coverage
 
     tracked_power_sensor = hass.data[DOMAIN][DATA_GROUP_ENTITIES]["sensor.tracked_power"]
-    assert tracked_power_sensor.entities == {"sensor.test1_power"}
+    assert tracked_power_sensor.entities == {"sensor.test1_power_new", "sensor.test3_power"}
     assert hass.states.get("sensor.tracked_power").state == "10.00"
 
     # Add a new power entity to registry
@@ -164,10 +173,10 @@ async def test_entity_registry_updates(hass: HomeAssistant) -> None:
         "sensor",
         "sensor",
         "aaa",
-        suggested_object_id="test3_power",
+        suggested_object_id="test4_power",
         original_device_class=SensorDeviceClass.POWER,
     )
 
     tracked_power_sensor = hass.data[DOMAIN][DATA_GROUP_ENTITIES]["sensor.tracked_power"]
-    assert tracked_power_sensor.entities == {"sensor.test1_power", "sensor.test3_power"}
+    assert tracked_power_sensor.entities == {"sensor.test1_power_new", "sensor.test3_power", "sensor.test4_power"}
     assert hass.states.get("sensor.tracked_power").state == "10.00"