From 0b5c533669d8138344dbe9bac6f9be330307c03f Mon Sep 17 00:00:00 2001 From: dougiteixeira <31328123+dougiteixeira@users.noreply.github.com> Date: Sat, 22 Jun 2024 10:31:47 -0300 Subject: [PATCH] Link the Trend helper entity to the source entity device (#119755) --- homeassistant/components/trend/__init__.py | 11 ++- .../components/trend/binary_sensor.py | 10 +++ tests/components/trend/test_binary_sensor.py | 45 ++++++++++ tests/components/trend/test_init.py | 87 ++++++++++++++++++- 4 files changed, 151 insertions(+), 2 deletions(-) diff --git a/homeassistant/components/trend/__init__.py b/homeassistant/components/trend/__init__.py index 7ec2d140c5e..c38730e7591 100644 --- a/homeassistant/components/trend/__init__.py +++ b/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)) diff --git a/homeassistant/components/trend/binary_sensor.py b/homeassistant/components/trend/binary_sensor.py index 2b70e2394f0..6788d22219b 100644 --- a/homeassistant/components/trend/binary_sensor.py +++ b/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 diff --git a/tests/components/trend/test_binary_sensor.py b/tests/components/trend/test_binary_sensor.py index 23d5a5357a7..ad85f65a9fc 100644 --- a/tests/components/trend/test_binary_sensor.py +++ b/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 diff --git a/tests/components/trend/test_init.py b/tests/components/trend/test_init.py index eea76025d65..7ffb18de297 100644 --- a/tests/components/trend/test_init.py +++ b/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