Add device link to History stats helper (#121848)

homeassistant/components/history_stats/__init__.py

@@ -7,6 +7,9 @@ from datetime import timedelta
 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_remove_stale_devices_links_keep_entity_device,
+)
 from homeassistant.helpers.template import Template
 
 from .const import CONF_DURATION, CONF_END, CONF_START, PLATFORMS
@@ -42,6 +45,12 @@ async def async_setup_entry(
     await coordinator.async_config_entry_first_refresh()
     entry.runtime_data = coordinator
 
+    async_remove_stale_devices_links_keep_entity_device(
+        hass,
+        entry.entry_id,
+        entry.options[CONF_ENTITY_ID],
+    )
+
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
     entry.async_on_unload(entry.add_update_listener(update_listener))
 

homeassistant/components/history_stats/sensor.py

@@ -26,6 +26,7 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import PlatformNotReady
 import homeassistant.helpers.config_validation as cv
+from homeassistant.helpers.device import async_device_info_to_link_from_entity
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.reload import async_setup_reload_service
 from homeassistant.helpers.template import Template
@@ -111,7 +112,9 @@ async def async_setup_platform(
     await coordinator.async_refresh()
     if not coordinator.last_update_success:
         raise PlatformNotReady from coordinator.last_exception
-    async_add_entities([HistoryStatsSensor(coordinator, sensor_type, name, unique_id)])
+    async_add_entities(
+        [HistoryStatsSensor(hass, coordinator, sensor_type, name, unique_id, entity_id)]
+    )
 
 
 async def async_setup_entry(
@@ -123,8 +126,13 @@ async def async_setup_entry(
 
     sensor_type: str = entry.options[CONF_TYPE]
     coordinator = entry.runtime_data
+    entity_id: str = entry.options[CONF_ENTITY_ID]
     async_add_entities(
-        [HistoryStatsSensor(coordinator, sensor_type, entry.title, entry.entry_id)]
+        [
+            HistoryStatsSensor(
+                hass, coordinator, sensor_type, entry.title, entry.entry_id, entity_id
+            )
+        ]
     )
 
 
@@ -167,16 +175,22 @@ class HistoryStatsSensor(HistoryStatsSensorBase):
 
     def __init__(
         self,
+        hass: HomeAssistant,
         coordinator: HistoryStatsUpdateCoordinator,
         sensor_type: str,
         name: str,
         unique_id: str | None,
+        source_entity_id: str,
     ) -> None:
         """Initialize the HistoryStats sensor."""
         super().__init__(coordinator, name)
         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(
+            hass,
+            source_entity_id,
+        )
         self._process_update()
         if self._type == CONF_TYPE_TIME:
             self._attr_device_class = SensorDeviceClass.DURATION

tests/components/history_stats/test_init.py

@@ -2,9 +2,17 @@
 
 from __future__ import annotations
 
+from homeassistant.components.history_stats.const import (
+    CONF_END,
+    CONF_START,
+    DEFAULT_NAME,
+    DOMAIN as HISTORY_STATS_DOMAIN,
+)
 from homeassistant.components.recorder import Recorder
 from homeassistant.config_entries import ConfigEntryState
+from homeassistant.const import CONF_ENTITY_ID, CONF_NAME, CONF_STATE, CONF_TYPE
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 
 from tests.common import MockConfigEntry
 
@@ -18,3 +26,93 @@ async def test_unload_entry(
     assert await hass.config_entries.async_unload(loaded_entry.entry_id)
     await hass.async_block_till_done()
     assert loaded_entry.state is ConfigEntryState.NOT_LOADED
+
+
+async def test_device_cleaning(
+    recorder_mock: Recorder,
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
+) -> None:
+    """Test the cleaning of devices linked to the helper History stats."""
+
+    # Source entity device config entry
+    source_config_entry = MockConfigEntry()
+    source_config_entry.add_to_hass(hass)
+
+    # Device entry of the source entity
+    source_device1_entry = device_registry.async_get_or_create(
+        config_entry_id=source_config_entry.entry_id,
+        identifiers={("binary_sensor", "identifier_test1")},
+        connections={("mac", "30:31:32:33:34:01")},
+    )
+
+    # Source entity registry
+    source_entity = entity_registry.async_get_or_create(
+        "binary_sensor",
+        "test",
+        "source",
+        config_entry=source_config_entry,
+        device_id=source_device1_entry.id,
+    )
+    await hass.async_block_till_done()
+    assert entity_registry.async_get("binary_sensor.test_source") is not None
+
+    # Configure the configuration entry for History stats
+    history_stats_config_entry = MockConfigEntry(
+        data={},
+        domain=HISTORY_STATS_DOMAIN,
+        options={
+            CONF_NAME: DEFAULT_NAME,
+            CONF_ENTITY_ID: "binary_sensor.test_source",
+            CONF_STATE: ["on"],
+            CONF_TYPE: "count",
+            CONF_START: "{{ as_timestamp(utcnow()) - 3600 }}",
+            CONF_END: "{{ utcnow() }}",
+        },
+        title="History stats",
+    )
+    history_stats_config_entry.add_to_hass(hass)
+    assert await hass.config_entries.async_setup(history_stats_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    # Confirm the link between the source entity device and the History stats sensor
+    history_stats_entity = entity_registry.async_get("sensor.history_stats")
+    assert history_stats_entity is not None
+    assert history_stats_entity.device_id == source_entity.device_id
+
+    # Device entry incorrectly linked to History stats config entry
+    device_registry.async_get_or_create(
+        config_entry_id=history_stats_config_entry.entry_id,
+        identifiers={("sensor", "identifier_test2")},
+        connections={("mac", "30:31:32:33:34:02")},
+    )
+    device_registry.async_get_or_create(
+        config_entry_id=history_stats_config_entry.entry_id,
+        identifiers={("sensor", "identifier_test3")},
+        connections={("mac", "30:31:32:33:34:03")},
+    )
+    await hass.async_block_till_done()
+
+    # Before reloading the config entry, two devices are expected to be linked
+    devices_before_reload = device_registry.devices.get_devices_for_config_entry_id(
+        history_stats_config_entry.entry_id
+    )
+    assert len(devices_before_reload) == 3
+
+    # Config entry reload
+    await hass.config_entries.async_reload(history_stats_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    # Confirm the link between the source entity device and the History stats sensor
+    history_stats_entity = entity_registry.async_get("sensor.history_stats")
+    assert history_stats_entity is not None
+    assert history_stats_entity.device_id == source_entity.device_id
+
+    # After reloading the config entry, only one linked device is expected
+    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 devices_after_reload[0].id == source_device1_entry.id

tests/components/history_stats/test_sensor.py

@@ -8,15 +8,28 @@ import pytest
 import voluptuous as vol
 
 from homeassistant import config as hass_config
-from homeassistant.components.history_stats.const import DOMAIN
+from homeassistant.components.history_stats.const import (
+    CONF_END,
+    CONF_START,
+    DEFAULT_NAME,
+    DOMAIN,
+)
 from homeassistant.components.history_stats.sensor import (
     PLATFORM_SCHEMA as SENSOR_SCHEMA,
 )
 from homeassistant.components.recorder import Recorder
-from homeassistant.const import ATTR_DEVICE_CLASS, SERVICE_RELOAD, STATE_UNKNOWN
+from homeassistant.const import (
+    ATTR_DEVICE_CLASS,
+    CONF_ENTITY_ID,
+    CONF_NAME,
+    CONF_STATE,
+    CONF_TYPE,
+    SERVICE_RELOAD,
+    STATE_UNKNOWN,
+)
 import homeassistant.core as ha
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry as er
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.helpers.entity_component import async_update_entity
 from homeassistant.setup import async_setup_component
 import homeassistant.util.dt as dt_util
@@ -1736,3 +1749,50 @@ async def test_unique_id(
         entity_registry.async_get("sensor.test").unique_id
         == "some_history_stats_unique_id"
     )
+
+
+async def test_device_id(
+    recorder_mock: Recorder,
+    hass: HomeAssistant,
+    entity_registry: er.EntityRegistry,
+    device_registry: dr.DeviceRegistry,
+) -> None:
+    """Test for source entity device for History stats."""
+    source_config_entry = MockConfigEntry()
+    source_config_entry.add_to_hass(hass)
+    source_device_entry = device_registry.async_get_or_create(
+        config_entry_id=source_config_entry.entry_id,
+        identifiers={("sensor", "identifier_test")},
+        connections={("mac", "30:31:32:33:34:35")},
+    )
+    source_entity = entity_registry.async_get_or_create(
+        "binary_sensor",
+        "test",
+        "source",
+        config_entry=source_config_entry,
+        device_id=source_device_entry.id,
+    )
+    await hass.async_block_till_done()
+    assert entity_registry.async_get("binary_sensor.test_source") is not None
+
+    history_stats_config_entry = MockConfigEntry(
+        data={},
+        domain=DOMAIN,
+        options={
+            CONF_NAME: DEFAULT_NAME,
+            CONF_ENTITY_ID: "binary_sensor.test_source",
+            CONF_STATE: ["on"],
+            CONF_TYPE: "count",
+            CONF_START: "{{ as_timestamp(utcnow()) - 3600 }}",
+            CONF_END: "{{ utcnow() }}",
+        },
+        title="History stats",
+    )
+    history_stats_config_entry.add_to_hass(hass)
+
+    assert await hass.config_entries.async_setup(history_stats_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    history_stats_entity = entity_registry.async_get("sensor.history_stats")
+    assert history_stats_entity is not None
+    assert history_stats_entity.device_id == source_entity.device_id