diff --git a/homeassistant/components/history_stats/__init__.py b/homeassistant/components/history_stats/__init__.py
index 476b17f581f647f8d953a8c5e8aa64432222c6fc..dcca10d73e9b96c865bc29198d724ea02baeefbd 100644
--- a/homeassistant/components/history_stats/__init__.py
+++ b/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))
diff --git a/homeassistant/components/history_stats/sensor.py b/homeassistant/components/history_stats/sensor.py
index a5139a8e9d68a243c5ed8fe4b5c3fe80139c678f..99e953ff9ddab263e34f52593a88141dff2b5ed1 100644
--- a/homeassistant/components/history_stats/sensor.py
+++ b/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
diff --git a/tests/components/history_stats/test_init.py b/tests/components/history_stats/test_init.py
index 180bb67e02a1a8373052ed74b984c904ef3535d0..4cd999ba31c6523f887de78f081d0e71e9bf19fc 100644
--- a/tests/components/history_stats/test_init.py
+++ b/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
diff --git a/tests/components/history_stats/test_sensor.py b/tests/components/history_stats/test_sensor.py
index 870c98503b47037b8bdd21139fa6f7732d07ca19..f86c04b3e5b044d79d002e1c875884f5844dedc5 100644
--- a/tests/components/history_stats/test_sensor.py
+++ b/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