diff --git a/homeassistant/components/statistics/__init__.py b/homeassistant/components/statistics/__init__.py index f800c82f1f9..34799e366d1 100644 --- a/homeassistant/components/statistics/__init__.py +++ b/homeassistant/components/statistics/__init__.py @@ -1,5 +1,7 @@ """The statistics component.""" +import logging + from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_ENTITY_ID, Platform from homeassistant.core import HomeAssistant @@ -7,15 +9,21 @@ 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.helper_integration import ( + async_handle_source_entity_changes, + async_remove_helper_config_entry_from_source_device, +) DOMAIN = "statistics" PLATFORMS = [Platform.SENSOR] +_LOGGER = logging.getLogger(__name__) + async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up Statistics from a config entry.""" + # This can be removed in HA Core 2026.2 async_remove_stale_devices_links_keep_entity_device( hass, entry.entry_id, @@ -36,6 +44,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: entry.async_on_unload( async_handle_source_entity_changes( hass, + add_helper_config_entry_to_device=False, 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( @@ -52,6 +61,40 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: return True +async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool: + """Migrate old entry.""" + _LOGGER.debug( + "Migrating from version %s.%s", config_entry.version, config_entry.minor_version + ) + + if config_entry.version > 1: + # This means the user has downgraded from a future version + return False + if config_entry.version == 1: + options = {**config_entry.options} + if config_entry.minor_version < 2: + # Remove the statistics config entry from the source device + if source_device_id := async_entity_id_to_device_id( + hass, options[CONF_ENTITY_ID] + ): + async_remove_helper_config_entry_from_source_device( + hass, + helper_config_entry_id=config_entry.entry_id, + source_device_id=source_device_id, + ) + hass.config_entries.async_update_entry( + config_entry, options=options, minor_version=2 + ) + + _LOGGER.debug( + "Migration to version %s.%s successful", + config_entry.version, + config_entry.minor_version, + ) + + return True + + async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Unload Statistics config entry.""" return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) diff --git a/homeassistant/components/statistics/config_flow.py b/homeassistant/components/statistics/config_flow.py index fb8c09868d5..d9ff172e0a4 100644 --- a/homeassistant/components/statistics/config_flow.py +++ b/homeassistant/components/statistics/config_flow.py @@ -161,6 +161,8 @@ OPTIONS_FLOW = { class StatisticsConfigFlowHandler(SchemaConfigFlowHandler, domain=DOMAIN): """Handle a config flow for Statistics.""" + MINOR_VERSION = 2 + config_flow = CONFIG_FLOW options_flow = OPTIONS_FLOW @@ -234,15 +236,15 @@ async def ws_start_preview( ) preview_entity = StatisticsSensor( hass, - entity_id, - name, - None, - state_characteristic, - sampling_size, - max_age, - msg["user_input"].get(CONF_KEEP_LAST_SAMPLE), - msg["user_input"].get(CONF_PRECISION), - msg["user_input"].get(CONF_PERCENTILE), + source_entity_id=entity_id, + name=name, + unique_id=None, + state_characteristic=state_characteristic, + samples_max_buffer_size=sampling_size, + samples_max_age=max_age, + samples_keep_last=msg["user_input"].get(CONF_KEEP_LAST_SAMPLE), + precision=msg["user_input"].get(CONF_PRECISION), + percentile=msg["user_input"].get(CONF_PERCENTILE), ) preview_entity.hass = hass diff --git a/homeassistant/components/statistics/sensor.py b/homeassistant/components/statistics/sensor.py index a5c5f10ecd0..8129a000b91 100644 --- a/homeassistant/components/statistics/sensor.py +++ b/homeassistant/components/statistics/sensor.py @@ -46,7 +46,7 @@ from homeassistant.core import ( split_entity_id, ) from homeassistant.helpers import config_validation as cv -from homeassistant.helpers.device import async_device_info_to_link_from_entity +from homeassistant.helpers.device import async_entity_id_to_device from homeassistant.helpers.entity_platform import ( AddConfigEntryEntitiesCallback, AddEntitiesCallback, @@ -659,6 +659,7 @@ class StatisticsSensor(SensorEntity): def __init__( self, hass: HomeAssistant, + *, source_entity_id: str, name: str, unique_id: str | None, @@ -673,10 +674,11 @@ class StatisticsSensor(SensorEntity): self._attr_name: str = name self._attr_unique_id: str | None = unique_id self._source_entity_id: str = source_entity_id - self._attr_device_info = async_device_info_to_link_from_entity( - hass, - source_entity_id, - ) + if source_entity_id: # Guard against empty source_entity_id in preview mode + self.device_entry = async_entity_id_to_device( + hass, + source_entity_id, + ) self.is_binary: bool = ( split_entity_id(self._source_entity_id)[0] == BINARY_SENSOR_DOMAIN ) diff --git a/tests/components/statistics/test_init.py b/tests/components/statistics/test_init.py index c11045a2eb2..2312daa8c52 100644 --- a/tests/components/statistics/test_init.py +++ b/tests/components/statistics/test_init.py @@ -10,7 +10,7 @@ from homeassistant.components import statistics from homeassistant.components.statistics import DOMAIN from homeassistant.components.statistics.config_flow import StatisticsConfigFlowHandler from homeassistant.config_entries import ConfigEntry, ConfigEntryState -from homeassistant.core import Event, HomeAssistant +from homeassistant.core import Event, HomeAssistant, callback from homeassistant.helpers import device_registry as dr, entity_registry as er from homeassistant.helpers.event import async_track_entity_registry_updated_event @@ -85,6 +85,7 @@ def track_entity_registry_actions(hass: HomeAssistant, entity_id: str) -> list[s """Track entity registry actions for an entity.""" events = [] + @callback def add_event(event: Event[er.EventEntityRegistryUpdatedData]) -> None: """Add entity registry updated event to the list.""" events.append(event.data["action"]) @@ -173,7 +174,7 @@ async def test_device_cleaning( devices_before_reload = device_registry.devices.get_devices_for_config_entry_id( statistics_config_entry.entry_id ) - assert len(devices_before_reload) == 3 + assert len(devices_before_reload) == 2 # Config entry reload await hass.config_entries.async_reload(statistics_config_entry.entry_id) @@ -188,9 +189,7 @@ async def test_device_cleaning( devices_after_reload = device_registry.devices.get_devices_for_config_entry_id( statistics_config_entry.entry_id ) - assert len(devices_after_reload) == 1 - - assert devices_after_reload[0].id == source_device1_entry.id + assert len(devices_after_reload) == 0 async def test_async_handle_source_entity_changes_source_entity_removed( @@ -201,6 +200,53 @@ async def test_async_handle_source_entity_changes_source_entity_removed( sensor_config_entry: ConfigEntry, sensor_device: dr.DeviceEntry, sensor_entity_entry: er.RegistryEntry, +) -> None: + """Test the statistics config entry is removed when the source entity is removed.""" + assert await hass.config_entries.async_setup(statistics_config_entry.entry_id) + await hass.async_block_till_done() + + statistics_entity_entry = entity_registry.async_get("sensor.my_statistics") + assert statistics_entity_entry.device_id == sensor_entity_entry.device_id + + sensor_device = device_registry.async_get(sensor_device.id) + assert statistics_config_entry.entry_id not in sensor_device.config_entries + + events = track_entity_registry_actions(hass, statistics_entity_entry.entity_id) + + # Remove the source sensor's config entry from the device, this removes the + # source sensor + with patch( + "homeassistant.components.statistics.async_unload_entry", + wraps=statistics.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 helper entity is removed + assert not entity_registry.async_get("sensor.my_statistics") + + # Check that the device is removed + assert not device_registry.async_get(sensor_device.id) + + # Check that the statistics config entry is removed + assert statistics_config_entry.entry_id not in hass.config_entries.async_entry_ids() + + # Check we got the expected events + assert events == ["remove"] + + +async def test_async_handle_source_entity_changes_source_entity_removed_shared_device( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, + statistics_config_entry: MockConfigEntry, + sensor_config_entry: ConfigEntry, + sensor_device: dr.DeviceEntry, + sensor_entity_entry: er.RegistryEntry, ) -> None: """Test the statistics config entry is removed when the source entity is removed.""" # Add another config entry to the sensor device @@ -217,7 +263,7 @@ async def test_async_handle_source_entity_changes_source_entity_removed( assert statistics_entity_entry.device_id == sensor_entity_entry.device_id sensor_device = device_registry.async_get(sensor_device.id) - assert statistics_config_entry.entry_id in sensor_device.config_entries + assert statistics_config_entry.entry_id not in sensor_device.config_entries events = track_entity_registry_actions(hass, statistics_entity_entry.entity_id) @@ -234,7 +280,10 @@ async def test_async_handle_source_entity_changes_source_entity_removed( await hass.async_block_till_done() mock_unload_entry.assert_called_once() - # Check that the statistics config entry is removed from the device + # Check that the helper entity is removed + assert not entity_registry.async_get("sensor.my_statistics") + + # Check that the statistics config entry is not in the device sensor_device = device_registry.async_get(sensor_device.id) assert statistics_config_entry.entry_id not in sensor_device.config_entries @@ -261,7 +310,7 @@ async def test_async_handle_source_entity_changes_source_entity_removed_from_dev assert statistics_entity_entry.device_id == sensor_entity_entry.device_id sensor_device = device_registry.async_get(sensor_device.id) - assert statistics_config_entry.entry_id in sensor_device.config_entries + assert statistics_config_entry.entry_id not in sensor_device.config_entries events = track_entity_registry_actions(hass, statistics_entity_entry.entity_id) @@ -276,7 +325,11 @@ async def test_async_handle_source_entity_changes_source_entity_removed_from_dev await hass.async_block_till_done() mock_unload_entry.assert_called_once() - # Check that the statistics config entry is removed from the device + # Check that the entity is no longer linked to the source device + statistics_entity_entry = entity_registry.async_get("sensor.my_statistics") + assert statistics_entity_entry.device_id is None + + # Check that the statistics config entry is not in the device sensor_device = device_registry.async_get(sensor_device.id) assert statistics_config_entry.entry_id not in sensor_device.config_entries @@ -309,7 +362,7 @@ async def test_async_handle_source_entity_changes_source_entity_moved_other_devi assert statistics_entity_entry.device_id == sensor_entity_entry.device_id sensor_device = device_registry.async_get(sensor_device.id) - assert statistics_config_entry.entry_id in sensor_device.config_entries + assert statistics_config_entry.entry_id not in sensor_device.config_entries sensor_device_2 = device_registry.async_get(sensor_device_2.id) assert statistics_config_entry.entry_id not in sensor_device_2.config_entries @@ -326,11 +379,15 @@ async def test_async_handle_source_entity_changes_source_entity_moved_other_devi await hass.async_block_till_done() mock_unload_entry.assert_called_once() - # Check that the statistics config entry is moved to the other device + # Check that the entity is linked to the other device + statistics_entity_entry = entity_registry.async_get("sensor.my_statistics") + assert statistics_entity_entry.device_id == sensor_device_2.id + + # Check that the history_stats config entry is not in any of the devices sensor_device = device_registry.async_get(sensor_device.id) assert statistics_config_entry.entry_id not in sensor_device.config_entries sensor_device_2 = device_registry.async_get(sensor_device_2.id) - assert statistics_config_entry.entry_id in sensor_device_2.config_entries + assert statistics_config_entry.entry_id not in sensor_device_2.config_entries # Check that the statistics config entry is not removed assert statistics_config_entry.entry_id in hass.config_entries.async_entry_ids() @@ -355,7 +412,7 @@ async def test_async_handle_source_entity_new_entity_id( assert statistics_entity_entry.device_id == sensor_entity_entry.device_id sensor_device = device_registry.async_get(sensor_device.id) - assert statistics_config_entry.entry_id in sensor_device.config_entries + assert statistics_config_entry.entry_id not in sensor_device.config_entries events = track_entity_registry_actions(hass, statistics_entity_entry.entity_id) @@ -373,12 +430,91 @@ async def test_async_handle_source_entity_new_entity_id( # Check that the statistics config entry is updated with the new entity ID assert statistics_config_entry.options["entity_id"] == "sensor.new_entity_id" - # Check that the helper config is still in the device + # Check that the helper config is not in the device sensor_device = device_registry.async_get(sensor_device.id) - assert statistics_config_entry.entry_id in sensor_device.config_entries + assert statistics_config_entry.entry_id not in sensor_device.config_entries # Check that the statistics config entry is not removed assert statistics_config_entry.entry_id in hass.config_entries.async_entry_ids() # Check we got the expected events assert events == [] + + +async def test_migration_1_1( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, + sensor_entity_entry: er.RegistryEntry, + sensor_device: dr.DeviceEntry, +) -> None: + """Test migration from v1.1 removes statistics config entry from device.""" + + statistics_config_entry = MockConfigEntry( + data={}, + domain=DOMAIN, + options={ + "name": "My statistics", + "entity_id": sensor_entity_entry.entity_id, + "state_characteristic": "mean", + "keep_last_sample": False, + "percentile": 50.0, + "precision": 2.0, + "sampling_size": 20.0, + }, + title="My statistics", + version=1, + minor_version=1, + ) + statistics_config_entry.add_to_hass(hass) + + # Add the helper config entry to the device + device_registry.async_update_device( + sensor_device.id, add_config_entry_id=statistics_config_entry.entry_id + ) + + # Check preconditions + sensor_device = device_registry.async_get(sensor_device.id) + assert statistics_config_entry.entry_id in sensor_device.config_entries + + await hass.config_entries.async_setup(statistics_config_entry.entry_id) + await hass.async_block_till_done() + + assert statistics_config_entry.state is ConfigEntryState.LOADED + + # Check that the helper config entry is removed from the device and the helper + # entity is linked to the source device + sensor_device = device_registry.async_get(sensor_device.id) + assert statistics_config_entry.entry_id not in sensor_device.config_entries + statistics_entity_entry = entity_registry.async_get("sensor.my_statistics") + assert statistics_entity_entry.device_id == sensor_entity_entry.device_id + + assert statistics_config_entry.version == 1 + assert statistics_config_entry.minor_version == 2 + + +async def test_migration_from_future_version( + hass: HomeAssistant, +) -> None: + """Test migration from future version.""" + config_entry = MockConfigEntry( + data={}, + domain=DOMAIN, + options={ + "name": "My statistics", + "entity_id": "sensor.test", + "state_characteristic": "mean", + "keep_last_sample": False, + "percentile": 50.0, + "precision": 2.0, + "sampling_size": 20.0, + }, + title="My statistics", + version=2, + minor_version=1, + ) + config_entry.add_to_hass(hass) + await hass.config_entries.async_setup(config_entry.entry_id) + await hass.async_block_till_done() + + assert config_entry.state is ConfigEntryState.MIGRATION_ERROR