fix: also remove member sensor from powercalc_groups storage when removed from group (#3972)

Author: bramstroker

custom_components/powercalc/sensors/group/custom.py

@@ -456,7 +456,6 @@ def __init__(
         self._attr_suggested_display_precision = self._rounding_digits
         if unique_id:
             self._attr_unique_id = unique_id
-        self._prev_state_store: PreviousStateStore = PreviousStateStore(hass)
         self._native_value_exact = Decimal(0)
         self._member_states: dict[str, Decimal] = {}
         self._ignore_unavailable_state = bool(self._sensor_config.get(CONF_IGNORE_UNAVAILABLE_STATE))
@@ -470,7 +469,6 @@ async def async_added_to_hass(self) -> None:
         """Register state listeners."""
         await super().async_added_to_hass()
 
-        self._prev_state_store = await PreviousStateStore.async_get_instance(self.hass)
         if self._update_interval > 0:
             self.async_on_remove(self._cancel_update_interval_exceeded_callback)
 
@@ -479,6 +477,12 @@ async def async_added_to_hass(self) -> None:
         if CONF_HIDE_MEMBERS in self._sensor_config:
             self._async_hide_members(bool(self._sensor_config.get(CONF_HIDE_MEMBERS)))
 
+        if not self._sensor_config.get(CONF_DISABLE_EXTENDED_ATTRIBUTES, False):
+            self._attr_extra_state_attributes = {
+                ATTR_ENTITIES: self._entities,
+                ATTR_IS_GROUP: True,
+            }
+
     async def async_will_remove_from_hass(self) -> None:
         """
         This will trigger when entity is about to be removed from HA
@@ -547,12 +551,6 @@ async def on_start(self, _: Any) -> None:  # noqa
             ),
         )
 
-        if not self._sensor_config.get(CONF_DISABLE_EXTENDED_ATTRIBUTES, False):
-            self._attr_extra_state_attributes = {
-                ATTR_ENTITIES: self._entities,
-                ATTR_IS_GROUP: True,
-            }
-
         await self.initial_update()
 
     async def initial_update(self) -> None:
@@ -720,9 +718,15 @@ def __init__(
             sensor_config.get(CONF_ENERGY_SENSOR_UNIT_PREFIX, UnitPrefix.NONE),
             UnitOfEnergy.WATT_HOUR,
         )
+        self._prev_state_store: PreviousStateStore = PreviousStateStore(hass)
 
     async def async_added_to_hass(self) -> None:
         """Register state listeners."""
+
+        self._prev_state_store = await PreviousStateStore.async_get_instance(self.hass)
+        # Clean up any entities that are no longer part of the group
+        self._prev_state_store.cleanup_entity_states(self.entity_id, self._entities)
+
         await self.restore_last_state()
 
         await super().async_added_to_hass()
@@ -900,6 +904,20 @@ def set_entity_state(self, group: str, entity_id: str, state: State) -> None:
         """Set the state for an energy sensor."""
         self.states.setdefault(group, {})[entity_id] = state
 
+    def cleanup_entity_states(self, group: str, current_entities: set[str]) -> None:
+        """Remove entity states that are no longer part of the group."""
+        group_states = self.states.get(group)
+        if group_states is None:
+            return
+
+        # Find entities that are in the store but not in the current set
+        entities_to_remove = set(group_states.keys()) - current_entities
+
+        # Remove those entities from the store
+        for entity_id in entities_to_remove:
+            _LOGGER.debug("Removing entity %s from group %s in PreviousStateStore", entity_id, group)
+            group_states.pop(entity_id, None)
+
     async def persist_states(self) -> None:
         """Save the current states to storage."""
         try:

tests/sensors/group/test_custom.py

@@ -1992,6 +1992,72 @@ async def test_resolve_entity_ids_skips_tasmota_yesterday_and_today(hass: HomeAs
     assert resolved == {"sensor.test_total"}
 
 
+async def test_remove_member_from_group(hass: HomeAssistant) -> None:
+    state_storage: PreviousStateStore = await PreviousStateStore.async_get_instance(hass)
+
+    # Setup 3 powercalc virtual power sensor and one group sensor
+    member_config_entries = []
+    member_entity_ids = []
+    for i in range(3):
+        config_entry = await setup_config_entry(
+            hass,
+            {
+                CONF_SENSOR_TYPE: SensorType.VIRTUAL_POWER,
+                CONF_UNIQUE_ID: f"pc{i}",
+                CONF_ENTITY_ID: DUMMY_ENTITY_ID,
+                CONF_NAME: f"VirtualSensor{i}",
+                CONF_MODE: CalculationStrategy.FIXED,
+                CONF_FIXED: {CONF_POWER: 50},
+            },
+        )
+        member_config_entries.append(config_entry)
+        member_entity_ids.append(f"sensor.virtualsensor{i}_energy")
+
+    group_entry = await setup_config_entry(
+        hass,
+        {
+            CONF_SENSOR_TYPE: SensorType.GROUP,
+            CONF_NAME: "TestGroup",
+            CONF_GROUP_MEMBER_SENSORS: [entry.entry_id for entry in member_config_entries],
+        },
+    )
+
+    # Trigger some changes on member power sensor so the group energy sensor integration logic is executed.
+    for i in range(3):
+        hass.states.async_set(f"sensor.virtualsensor{i}_power", "60.00")
+        await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    # Assert the group sensor has the correct 3 member sensors added
+    group_state = hass.states.get("sensor.testgroup_energy")
+    assert group_state
+    assert group_state.attributes.get(CONF_ENTITIES) == set(member_entity_ids)
+
+    assert state_storage.get_entity_state("sensor.testgroup_energy", member_entity_ids[0])
+    assert state_storage.get_entity_state("sensor.testgroup_energy", member_entity_ids[1])
+    assert state_storage.get_entity_state("sensor.testgroup_energy", member_entity_ids[2])
+
+    # Remove one powercalc sensor from the group
+    member_config_entries.remove(member_config_entries[1])
+    member_entity_ids.remove(member_entity_ids[1])
+    hass.config_entries.async_update_entry(
+        group_entry,
+        data={
+            **group_entry.data,
+            CONF_GROUP_MEMBER_SENSORS: [entry.entry_id for entry in member_config_entries],
+        },
+    )
+
+    # Assert the group sensor has the correct 2 member sensors added, and the removed sensor is not there anymore
+    group_state = hass.states.get("sensor.testgroup_energy")
+    assert group_state
+    assert group_state.attributes.get(CONF_ENTITIES) == set(member_entity_ids)
+
+    assert state_storage.get_entity_state("sensor.testgroup_energy", member_entity_ids[0])
+    assert not state_storage.get_entity_state("sensor.testgroup_energy", "sensor.virtualsensor1_energy")
+    assert state_storage.get_entity_state("sensor.testgroup_energy", member_entity_ids[1])
+
+
 async def _create_energy_group(
     hass: HomeAssistant,
     name: str,