Do not add history_stats config entry to source device (#148729)

homeassistant/components/history_stats/__init__.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from datetime import timedelta
+import logging
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_ENTITY_ID, CONF_STATE
@@ -11,7 +12,10 @@ 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,
+)
 from homeassistant.helpers.template import Template
 
 from .const import CONF_DURATION, CONF_END, CONF_START, PLATFORMS
@@ -20,6 +24,8 @@ from .data import HistoryStats
 
 type HistoryStatsConfigEntry = ConfigEntry[HistoryStatsUpdateCoordinator]
 
+_LOGGER = logging.getLogger(__name__)
+
 
 async def async_setup_entry(
     hass: HomeAssistant, entry: HistoryStatsConfigEntry
@@ -47,6 +53,7 @@ async def async_setup_entry(
     await coordinator.async_config_entry_first_refresh()
     entry.runtime_data = coordinator
 
+    # This can be removed in HA Core 2026.2
     async_remove_stale_devices_links_keep_entity_device(
         hass,
         entry.entry_id,
@@ -67,6 +74,7 @@ async def async_setup_entry(
     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(
@@ -83,6 +91,40 @@ async def async_setup_entry(
     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 history_stats 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: HistoryStatsConfigEntry
 ) -> bool:

homeassistant/components/history_stats/config_flow.py

@@ -124,6 +124,8 @@ OPTIONS_FLOW = {
 class HistoryStatsConfigFlowHandler(SchemaConfigFlowHandler, domain=DOMAIN):
     """Handle a config flow for History stats."""
 
+    MINOR_VERSION = 2
+
     config_flow = CONFIG_FLOW
     options_flow = OPTIONS_FLOW
 
@@ -229,7 +231,12 @@ async def ws_start_preview(
     coordinator = HistoryStatsUpdateCoordinator(hass, history_stats, None, name, True)
     await coordinator.async_refresh()
     preview_entity = HistoryStatsSensor(
-        hass, coordinator, sensor_type, name, None, entity_id
+        hass,
+        coordinator=coordinator,
+        sensor_type=sensor_type,
+        name=name,
+        unique_id=None,
+        source_entity_id=entity_id,
     )
     preview_entity.hass = hass
 

homeassistant/components/history_stats/sensor.py

@@ -27,7 +27,7 @@ from homeassistant.const import (
 from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
 from homeassistant.exceptions import PlatformNotReady
 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,
@@ -113,7 +113,16 @@ async def async_setup_platform(
     if not coordinator.last_update_success:
         raise PlatformNotReady from coordinator.last_exception
     async_add_entities(
-        [HistoryStatsSensor(hass, coordinator, sensor_type, name, unique_id, entity_id)]
+        [
+            HistoryStatsSensor(
+                hass,
+                coordinator=coordinator,
+                sensor_type=sensor_type,
+                name=name,
+                unique_id=unique_id,
+                source_entity_id=entity_id,
+            )
+        ]
     )
 
 
@@ -130,7 +139,12 @@ async def async_setup_entry(
     async_add_entities(
         [
             HistoryStatsSensor(
-                hass, coordinator, sensor_type, entry.title, entry.entry_id, entity_id
+                hass,
+                coordinator=coordinator,
+                sensor_type=sensor_type,
+                name=entry.title,
+                unique_id=entry.entry_id,
+                source_entity_id=entity_id,
             )
         ]
     )
@@ -176,6 +190,7 @@ class HistoryStatsSensor(HistoryStatsSensorBase):
     def __init__(
         self,
         hass: HomeAssistant,
+        *,
         coordinator: HistoryStatsUpdateCoordinator,
         sensor_type: str,
         name: str,
@@ -190,7 +205,8 @@ class HistoryStatsSensor(HistoryStatsSensorBase):
         self._attr_native_unit_of_measurement = UNITS[sensor_type]
         self._type = sensor_type
         self._attr_unique_id = unique_id
-        self._attr_device_info = async_device_info_to_link_from_entity(
+        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,
             )

tests/components/history_stats/test_init.py

@@ -18,7 +18,7 @@ from homeassistant.components.history_stats.const import (
 )
 from homeassistant.config_entries import ConfigEntry, ConfigEntryState
 from homeassistant.const import CONF_ENTITY_ID, CONF_NAME, CONF_STATE, CONF_TYPE
-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
 
@@ -92,6 +92,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"])
@@ -181,7 +182,7 @@ async def test_device_cleaning(
     devices_before_reload = device_registry.devices.get_devices_for_config_entry_id(
         history_stats_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(history_stats_config_entry.entry_id)
@@ -196,9 +197,7 @@ async def test_device_cleaning(
     devices_after_reload = device_registry.devices.get_devices_for_config_entry_id(
         history_stats_config_entry.entry_id
     )
-    assert len(devices_after_reload) == 1
+    assert len(devices_after_reload) == 0
-
-    assert devices_after_reload[0].id == source_device1_entry.id
 
 
 @pytest.mark.usefixtures("recorder_mock")
@@ -210,6 +209,56 @@ 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 history_stats config entry is removed when the source entity is removed."""
+    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 not 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 helper entity is removed
+    assert not entity_registry.async_get("sensor.my_history_stats")
+
+    # Check that the device is removed
+    assert not device_registry.async_get(sensor_device.id)
+
+    # 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_shared_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 history_stats config entry is removed when the source entity is removed."""
     # Add another config entry to the sensor device
@@ -226,7 +275,7 @@ async def test_async_handle_source_entity_changes_source_entity_removed(
     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
+    assert history_stats_config_entry.entry_id not in sensor_device.config_entries
 
     events = track_entity_registry_actions(hass, history_stats_entity_entry.entity_id)
 
@@ -243,7 +292,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 history_stats config entry is removed from the device
+    # Check that the helper entity is removed
+    assert not entity_registry.async_get("sensor.my_history_stats")
+
+    # Check that the history_stats config entry is not in the device
     sensor_device = device_registry.async_get(sensor_device.id)
     assert history_stats_config_entry.entry_id not in sensor_device.config_entries
 
@@ -273,7 +325,7 @@ async def test_async_handle_source_entity_changes_source_entity_removed_from_dev
     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
+    assert history_stats_config_entry.entry_id not in sensor_device.config_entries
 
     events = track_entity_registry_actions(hass, history_stats_entity_entry.entity_id)
 
@@ -288,7 +340,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 history_stats config entry is removed from the device
+    # Check that the entity is no longer linked to the source device
+    history_stats_entity_entry = entity_registry.async_get("sensor.my_history_stats")
+    assert history_stats_entity_entry.device_id is None
+
+    # Check that the history_stats config entry is not in the device
     sensor_device = device_registry.async_get(sensor_device.id)
     assert history_stats_config_entry.entry_id not in sensor_device.config_entries
 
@@ -322,7 +378,7 @@ async def test_async_handle_source_entity_changes_source_entity_moved_other_devi
     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
+    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 not in sensor_device_2.config_entries
 
@@ -339,11 +395,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 history_stats config entry is moved to the other device
+    # Check that the entity is linked to the other device
+    history_stats_entity_entry = entity_registry.async_get("sensor.my_history_stats")
+    assert history_stats_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 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
+    assert history_stats_config_entry.entry_id not 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()
@@ -369,7 +429,7 @@ async def test_async_handle_source_entity_new_entity_id(
     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
+    assert history_stats_config_entry.entry_id not in sensor_device.config_entries
 
     events = track_entity_registry_actions(hass, history_stats_entity_entry.entity_id)
 
@@ -387,12 +447,91 @@ async def test_async_handle_source_entity_new_entity_id(
     # 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
+    # Check that the helper config is not in the device
     sensor_device = device_registry.async_get(sensor_device.id)
-    assert history_stats_config_entry.entry_id in sensor_device.config_entries
+    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 == []
+
+
+@pytest.mark.usefixtures("recorder_mock")
+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 history_stats config entry from device."""
+
+    history_stats_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=1,
+        minor_version=1,
+    )
+    history_stats_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=history_stats_config_entry.entry_id
+    )
+
+    # Check preconditions
+    sensor_device = device_registry.async_get(sensor_device.id)
+    assert history_stats_config_entry.entry_id in sensor_device.config_entries
+
+    await hass.config_entries.async_setup(history_stats_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert history_stats_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 history_stats_config_entry.entry_id not in sensor_device.config_entries
+    history_stats_entity_entry = entity_registry.async_get("sensor.my_history_stats")
+    assert history_stats_entity_entry.device_id == sensor_entity_entry.device_id
+
+    assert history_stats_config_entry.version == 1
+    assert history_stats_config_entry.minor_version == 2
+
+
+@pytest.mark.usefixtures("recorder_mock")
+async def test_migration_from_future_version(
+    hass: HomeAssistant,
+) -> None:
+    """Test migration from future version."""
+    config_entry = MockConfigEntry(
+        data={},
+        domain=DOMAIN,
+        options={
+            CONF_NAME: DEFAULT_NAME,
+            CONF_ENTITY_ID: "sensor.test",
+            CONF_STATE: ["on"],
+            CONF_TYPE: "count",
+            CONF_START: "{{ as_timestamp(utcnow()) - 3600 }}",
+            CONF_END: "{{ utcnow() }}",
+        },
+        title="My history stats",
+        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