Link the Trend helper entity to the source entity device (#119755)

homeassistant/components/trend/__init__.py

@@ -3,8 +3,11 @@
 from __future__ import annotations
 
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import Platform
+from homeassistant.const import CONF_ENTITY_ID, Platform
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers.device import (
+    async_remove_stale_devices_links_keep_entity_device,
+)
 
 PLATFORMS = [Platform.BINARY_SENSOR]
 
@@ -12,6 +15,12 @@ PLATFORMS = [Platform.BINARY_SENSOR]
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Trend from a config entry."""
 
+    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(config_entry_update_listener))
 

homeassistant/components/trend/binary_sensor.py

@@ -32,7 +32,9 @@ from homeassistant.const import (
     STATE_UNKNOWN,
 )
 from homeassistant.core import Event, EventStateChangedData, HomeAssistant, callback
+from homeassistant.helpers import device_registry as dr
 import homeassistant.helpers.config_validation as cv
+from homeassistant.helpers.device import async_device_info_to_link_from_entity
 from homeassistant.helpers.entity import generate_entity_id
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.event import async_track_state_change_event
@@ -133,6 +135,11 @@ async def async_setup_entry(
 ) -> None:
     """Set up trend sensor from config entry."""
 
+    device_info = async_device_info_to_link_from_entity(
+        hass,
+        entry.options[CONF_ENTITY_ID],
+    )
+
     async_add_entities(
         [
             SensorTrend(
@@ -147,6 +154,7 @@ async def async_setup_entry(
                 min_samples=entry.options.get(CONF_MIN_SAMPLES, DEFAULT_MIN_SAMPLES),
                 max_samples=entry.options.get(CONF_MAX_SAMPLES, DEFAULT_MAX_SAMPLES),
                 unique_id=entry.entry_id,
+                device_info=device_info,
             )
         ]
     )
@@ -172,6 +180,7 @@ class SensorTrend(BinarySensorEntity, RestoreEntity):
         unique_id: str | None = None,
         device_class: BinarySensorDeviceClass | None = None,
         sensor_entity_id: str | None = None,
+        device_info: dr.DeviceInfo | None = None,
     ) -> None:
         """Initialize the sensor."""
         self._entity_id = entity_id
@@ -185,6 +194,7 @@ class SensorTrend(BinarySensorEntity, RestoreEntity):
         self._attr_name = name
         self._attr_device_class = device_class
         self._attr_unique_id = unique_id
+        self._attr_device_info = device_info
 
         if sensor_entity_id:
             self.entity_id = sensor_entity_id

tests/components/trend/test_binary_sensor.py

@@ -8,8 +8,10 @@ from freezegun.api import FrozenDateTimeFactory
 import pytest
 
 from homeassistant import setup
+from homeassistant.components.trend.const import DOMAIN
 from homeassistant.const import STATE_OFF, STATE_ON, STATE_UNKNOWN
 from homeassistant.core import HomeAssistant, State
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.setup import async_setup_component
 
 from .conftest import ComponentSetup
@@ -350,3 +352,46 @@ async def test_invalid_min_sample(
         "Invalid config for 'binary_sensor' from integration 'trend': min_samples must "
         "be smaller than or equal to max_samples" in record.message
     )
+
+
+async def test_device_id(
+    hass: HomeAssistant,
+    entity_registry: er.EntityRegistry,
+    device_registry: dr.DeviceRegistry,
+) -> None:
+    """Test for source entity device for Trend."""
+    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(
+        "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("sensor.test_source") is not None
+
+    trend_config_entry = MockConfigEntry(
+        data={},
+        domain=DOMAIN,
+        options={
+            "name": "Trend",
+            "entity_id": "sensor.test_source",
+            "invert": False,
+        },
+        title="Trend",
+    )
+    trend_config_entry.add_to_hass(hass)
+
+    assert await hass.config_entries.async_setup(trend_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    trend_entity = entity_registry.async_get("binary_sensor.trend")
+    assert trend_entity is not None
+    assert trend_entity.device_id == source_entity.device_id

tests/components/trend/test_init.py

@@ -1,8 +1,9 @@
 """Test the Trend integration."""
 
+from homeassistant.components.trend.const import DOMAIN
 from homeassistant.config_entries import ConfigEntryState
 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 .conftest import ComponentSetup
 
@@ -50,3 +51,87 @@ async def test_reload_config_entry(
 
     assert config_entry.state is ConfigEntryState.LOADED
     assert config_entry.data == {**config_entry.data, "max_samples": 4.0}
+
+
+async def test_device_cleaning(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
+) -> None:
+    """Test for source entity device for Trend."""
+
+    # 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={("sensor", "identifier_test1")},
+        connections={("mac", "30:31:32:33:34:01")},
+    )
+
+    # Source entity registry
+    source_entity = entity_registry.async_get_or_create(
+        "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("sensor.test_source") is not None
+
+    # Configure the configuration entry for Trend
+    trend_config_entry = MockConfigEntry(
+        data={},
+        domain=DOMAIN,
+        options={
+            "name": "Trend",
+            "entity_id": "sensor.test_source",
+            "invert": False,
+        },
+        title="Trend",
+    )
+    trend_config_entry.add_to_hass(hass)
+    assert await hass.config_entries.async_setup(trend_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    # Confirm the link between the source entity device and the trend sensor
+    trend_entity = entity_registry.async_get("binary_sensor.trend")
+    assert trend_entity is not None
+    assert trend_entity.device_id == source_entity.device_id
+
+    # Device entry incorrectly linked to Trend config entry
+    device_registry.async_get_or_create(
+        config_entry_id=trend_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=trend_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(
+        trend_config_entry.entry_id
+    )
+    assert len(devices_before_reload) == 3
+
+    # Config entry reload
+    await hass.config_entries.async_reload(trend_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    # Confirm the link between the source entity device and the trend sensor after reload
+    trend_entity = entity_registry.async_get("binary_sensor.trend")
+    assert trend_entity is not None
+    assert trend_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(
+        trend_config_entry.entry_id
+    )
+    assert len(devices_after_reload) == 1