From cf5cdf3cdb0f1d7b2bdb214a0906b3121d60a6b7 Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Wed, 11 Jun 2025 15:16:51 +0200 Subject: [PATCH] Handle changes to source entity in history_stats helper (#146521) --- .../components/history_stats/__init__.py | 26 ++ .../components/history_stats/config_flow.py | 2 +- tests/components/history_stats/test_init.py | 294 +++++++++++++++++- 3 files changed, 314 insertions(+), 8 deletions(-) diff --git a/homeassistant/components/history_stats/__init__.py b/homeassistant/components/history_stats/__init__.py index 63f32138dba..a3565f9ed77 100644 --- a/homeassistant/components/history_stats/__init__.py +++ b/homeassistant/components/history_stats/__init__.py @@ -8,8 +8,10 @@ from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_ENTITY_ID, CONF_STATE from homeassistant.core import HomeAssistant from homeassistant.helpers.device import ( + async_entity_id_to_device_id, async_remove_stale_devices_links_keep_entity_device, ) +from homeassistant.helpers.helper_integration import async_handle_source_entity_changes from homeassistant.helpers.template import Template from .const import CONF_DURATION, CONF_END, CONF_START, PLATFORMS @@ -51,6 +53,30 @@ async def async_setup_entry( entry.options[CONF_ENTITY_ID], ) + def set_source_entity_id_or_uuid(source_entity_id: str) -> None: + hass.config_entries.async_update_entry( + entry, + options={**entry.options, CONF_ENTITY_ID: source_entity_id}, + ) + + async def source_entity_removed() -> None: + # The source entity has been removed, we remove the config entry because + # history_stats does not allow replacing the input entity. + await hass.config_entries.async_remove(entry.entry_id) + + entry.async_on_unload( + async_handle_source_entity_changes( + hass, + helper_config_entry_id=entry.entry_id, + set_source_entity_id_or_uuid=set_source_entity_id_or_uuid, + source_device_id=async_entity_id_to_device_id( + hass, entry.options[CONF_ENTITY_ID] + ), + source_entity_id_or_uuid=entry.options[CONF_ENTITY_ID], + source_entity_removed=source_entity_removed, + ) + ) + await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) entry.async_on_unload(entry.add_update_listener(update_listener)) diff --git a/homeassistant/components/history_stats/config_flow.py b/homeassistant/components/history_stats/config_flow.py index 96c8f319fbc..ca3d5229b6b 100644 --- a/homeassistant/components/history_stats/config_flow.py +++ b/homeassistant/components/history_stats/config_flow.py @@ -107,7 +107,7 @@ OPTIONS_FLOW = { } -class StatisticsConfigFlowHandler(SchemaConfigFlowHandler, domain=DOMAIN): +class HistoryStatsConfigFlowHandler(SchemaConfigFlowHandler, domain=DOMAIN): """Handle a config flow for History stats.""" config_flow = CONFIG_FLOW diff --git a/tests/components/history_stats/test_init.py b/tests/components/history_stats/test_init.py index c99d836a822..cb3350f497f 100644 --- a/tests/components/history_stats/test_init.py +++ b/tests/components/history_stats/test_init.py @@ -2,24 +2,107 @@ from __future__ import annotations +from unittest.mock import patch + +import pytest + +from homeassistant.components import history_stats +from homeassistant.components.history_stats.config_flow import ( + HistoryStatsConfigFlowHandler, +) from homeassistant.components.history_stats.const import ( CONF_END, CONF_START, DEFAULT_NAME, DOMAIN, ) -from homeassistant.components.recorder import Recorder -from homeassistant.config_entries import ConfigEntryState +from homeassistant.config_entries import ConfigEntry, ConfigEntryState from homeassistant.const import CONF_ENTITY_ID, CONF_NAME, CONF_STATE, CONF_TYPE -from homeassistant.core import HomeAssistant +from homeassistant.core import Event, HomeAssistant from homeassistant.helpers import device_registry as dr, entity_registry as er +from homeassistant.helpers.event import async_track_entity_registry_updated_event from tests.common import MockConfigEntry -async def test_unload_entry( - recorder_mock: Recorder, hass: HomeAssistant, loaded_entry: MockConfigEntry -) -> None: +@pytest.fixture +def sensor_config_entry(hass: HomeAssistant) -> er.RegistryEntry: + """Fixture to create a sensor config entry.""" + sensor_config_entry = MockConfigEntry() + sensor_config_entry.add_to_hass(hass) + return sensor_config_entry + + +@pytest.fixture +def sensor_device( + device_registry: dr.DeviceRegistry, sensor_config_entry: ConfigEntry +) -> dr.DeviceEntry: + """Fixture to create a sensor device.""" + return device_registry.async_get_or_create( + config_entry_id=sensor_config_entry.entry_id, + connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")}, + ) + + +@pytest.fixture +def sensor_entity_entry( + entity_registry: er.EntityRegistry, + sensor_config_entry: ConfigEntry, + sensor_device: dr.DeviceEntry, +) -> er.RegistryEntry: + """Fixture to create a sensor entity entry.""" + return entity_registry.async_get_or_create( + "sensor", + "test", + "unique", + config_entry=sensor_config_entry, + device_id=sensor_device.id, + original_name="ABC", + ) + + +@pytest.fixture +def history_stats_config_entry( + hass: HomeAssistant, + sensor_entity_entry: er.RegistryEntry, +) -> MockConfigEntry: + """Fixture to create a history_stats config entry.""" + config_entry = MockConfigEntry( + data={}, + domain=DOMAIN, + options={ + CONF_NAME: DEFAULT_NAME, + CONF_ENTITY_ID: sensor_entity_entry.entity_id, + CONF_STATE: ["on"], + CONF_TYPE: "count", + CONF_START: "{{ as_timestamp(utcnow()) - 3600 }}", + CONF_END: "{{ utcnow() }}", + }, + title="My history stats", + version=HistoryStatsConfigFlowHandler.VERSION, + minor_version=HistoryStatsConfigFlowHandler.MINOR_VERSION, + ) + + config_entry.add_to_hass(hass) + + return config_entry + + +def track_entity_registry_actions(hass: HomeAssistant, entity_id: str) -> list[str]: + """Track entity registry actions for an entity.""" + events = [] + + def add_event(event: Event[er.EventEntityRegistryUpdatedData]) -> None: + """Add entity registry updated event to the list.""" + events.append(event.data["action"]) + + async_track_entity_registry_updated_event(hass, entity_id, add_event) + + return events + + +@pytest.mark.usefixtures("recorder_mock") +async def test_unload_entry(hass: HomeAssistant, loaded_entry: MockConfigEntry) -> None: """Test unload an entry.""" assert loaded_entry.state is ConfigEntryState.LOADED @@ -28,8 +111,8 @@ async def test_unload_entry( assert loaded_entry.state is ConfigEntryState.NOT_LOADED +@pytest.mark.usefixtures("recorder_mock") async def test_device_cleaning( - recorder_mock: Recorder, hass: HomeAssistant, device_registry: dr.DeviceRegistry, entity_registry: er.EntityRegistry, @@ -116,3 +199,200 @@ async def test_device_cleaning( assert len(devices_after_reload) == 1 assert devices_after_reload[0].id == source_device1_entry.id + + +@pytest.mark.usefixtures("recorder_mock") +async def test_async_handle_source_entity_changes_source_entity_removed( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, + history_stats_config_entry: MockConfigEntry, + sensor_config_entry: ConfigEntry, + sensor_device: dr.DeviceEntry, + sensor_entity_entry: er.RegistryEntry, +) -> None: + """Test the history_stats config entry is removed when the source entity is removed.""" + # Add another config entry to the sensor device + other_config_entry = MockConfigEntry() + other_config_entry.add_to_hass(hass) + device_registry.async_update_device( + sensor_device.id, add_config_entry_id=other_config_entry.entry_id + ) + + assert await hass.config_entries.async_setup(history_stats_config_entry.entry_id) + await hass.async_block_till_done() + + history_stats_entity_entry = entity_registry.async_get("sensor.my_history_stats") + assert history_stats_entity_entry.device_id == sensor_entity_entry.device_id + + sensor_device = device_registry.async_get(sensor_device.id) + assert history_stats_config_entry.entry_id in sensor_device.config_entries + + events = track_entity_registry_actions(hass, history_stats_entity_entry.entity_id) + + # Remove the source sensor's config entry from the device, this removes the + # source sensor + with patch( + "homeassistant.components.history_stats.async_unload_entry", + wraps=history_stats.async_unload_entry, + ) as mock_unload_entry: + device_registry.async_update_device( + sensor_device.id, remove_config_entry_id=sensor_config_entry.entry_id + ) + await hass.async_block_till_done() + await hass.async_block_till_done() + mock_unload_entry.assert_called_once() + + # Check that the history_stats config entry is removed from the device + sensor_device = device_registry.async_get(sensor_device.id) + assert history_stats_config_entry.entry_id not in sensor_device.config_entries + + # Check that the history_stats config entry is removed + assert ( + history_stats_config_entry.entry_id not in hass.config_entries.async_entry_ids() + ) + + # Check we got the expected events + assert events == ["remove"] + + +@pytest.mark.usefixtures("recorder_mock") +async def test_async_handle_source_entity_changes_source_entity_removed_from_device( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, + history_stats_config_entry: MockConfigEntry, + sensor_device: dr.DeviceEntry, + sensor_entity_entry: er.RegistryEntry, +) -> None: + """Test the source entity removed from the source device.""" + assert await hass.config_entries.async_setup(history_stats_config_entry.entry_id) + await hass.async_block_till_done() + + history_stats_entity_entry = entity_registry.async_get("sensor.my_history_stats") + assert history_stats_entity_entry.device_id == sensor_entity_entry.device_id + + sensor_device = device_registry.async_get(sensor_device.id) + assert history_stats_config_entry.entry_id in sensor_device.config_entries + + events = track_entity_registry_actions(hass, history_stats_entity_entry.entity_id) + + # Remove the source sensor from the device + with patch( + "homeassistant.components.history_stats.async_unload_entry", + wraps=history_stats.async_unload_entry, + ) as mock_unload_entry: + entity_registry.async_update_entity( + sensor_entity_entry.entity_id, device_id=None + ) + await hass.async_block_till_done() + mock_unload_entry.assert_called_once() + + # Check that the history_stats config entry is removed from the device + sensor_device = device_registry.async_get(sensor_device.id) + assert history_stats_config_entry.entry_id not in sensor_device.config_entries + + # Check that the history_stats config entry is not removed + assert history_stats_config_entry.entry_id in hass.config_entries.async_entry_ids() + + # Check we got the expected events + assert events == ["update"] + + +@pytest.mark.usefixtures("recorder_mock") +async def test_async_handle_source_entity_changes_source_entity_moved_other_device( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, + history_stats_config_entry: MockConfigEntry, + sensor_config_entry: ConfigEntry, + sensor_device: dr.DeviceEntry, + sensor_entity_entry: er.RegistryEntry, +) -> None: + """Test the source entity is moved to another device.""" + sensor_device_2 = device_registry.async_get_or_create( + config_entry_id=sensor_config_entry.entry_id, + connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:FF")}, + ) + + assert await hass.config_entries.async_setup(history_stats_config_entry.entry_id) + await hass.async_block_till_done() + + history_stats_entity_entry = entity_registry.async_get("sensor.my_history_stats") + assert history_stats_entity_entry.device_id == sensor_entity_entry.device_id + + sensor_device = device_registry.async_get(sensor_device.id) + assert history_stats_config_entry.entry_id in sensor_device.config_entries + sensor_device_2 = device_registry.async_get(sensor_device_2.id) + assert history_stats_config_entry.entry_id not in sensor_device_2.config_entries + + events = track_entity_registry_actions(hass, history_stats_entity_entry.entity_id) + + # Move the source sensor to another device + with patch( + "homeassistant.components.history_stats.async_unload_entry", + wraps=history_stats.async_unload_entry, + ) as mock_unload_entry: + entity_registry.async_update_entity( + sensor_entity_entry.entity_id, device_id=sensor_device_2.id + ) + await hass.async_block_till_done() + mock_unload_entry.assert_called_once() + + # Check that the history_stats config entry is moved to the other device + sensor_device = device_registry.async_get(sensor_device.id) + assert history_stats_config_entry.entry_id not in sensor_device.config_entries + sensor_device_2 = device_registry.async_get(sensor_device_2.id) + assert history_stats_config_entry.entry_id in sensor_device_2.config_entries + + # Check that the history_stats config entry is not removed + assert history_stats_config_entry.entry_id in hass.config_entries.async_entry_ids() + + # Check we got the expected events + assert events == ["update"] + + +@pytest.mark.usefixtures("recorder_mock") +async def test_async_handle_source_entity_new_entity_id( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, + history_stats_config_entry: MockConfigEntry, + sensor_device: dr.DeviceEntry, + sensor_entity_entry: er.RegistryEntry, +) -> None: + """Test the source entity's entity ID is changed.""" + assert await hass.config_entries.async_setup(history_stats_config_entry.entry_id) + await hass.async_block_till_done() + + history_stats_entity_entry = entity_registry.async_get("sensor.my_history_stats") + assert history_stats_entity_entry.device_id == sensor_entity_entry.device_id + + sensor_device = device_registry.async_get(sensor_device.id) + assert history_stats_config_entry.entry_id in sensor_device.config_entries + + events = track_entity_registry_actions(hass, history_stats_entity_entry.entity_id) + + # Change the source entity's entity ID + with patch( + "homeassistant.components.history_stats.async_unload_entry", + wraps=history_stats.async_unload_entry, + ) as mock_unload_entry: + entity_registry.async_update_entity( + sensor_entity_entry.entity_id, new_entity_id="sensor.new_entity_id" + ) + await hass.async_block_till_done() + mock_unload_entry.assert_called_once() + + # Check that the history_stats config entry is updated with the new entity ID + assert history_stats_config_entry.options[CONF_ENTITY_ID] == "sensor.new_entity_id" + + # Check that the helper config is still in the device + sensor_device = device_registry.async_get(sensor_device.id) + assert history_stats_config_entry.entry_id in sensor_device.config_entries + + # Check that the history_stats config entry is not removed + assert history_stats_config_entry.entry_id in hass.config_entries.async_entry_ids() + + # Check we got the expected events + assert events == []