Skip to content

Commit

Permalink
fix: fully implement entity registry updating and 100% coverage
Browse files Browse the repository at this point in the history
  • Loading branch information
bramstroker committed Dec 26, 2024
1 parent 4aad8cb commit 6faa47c
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 12 deletions.
19 changes: 13 additions & 6 deletions custom_components/powercalc/sensors/group/tracked_untracked.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
21 changes: 15 additions & 6 deletions tests/sensors/group/test_tracked_untracked.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -146,28 +150,33 @@ 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
entity_registry.async_get_or_create(
"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"

0 comments on commit 6faa47c

Please sign in to comment.