From 8ed11d4b90966de616de44bfe39add1162d9893c Mon Sep 17 00:00:00 2001 From: dougiteixeira <31328123+dougiteixeira@users.noreply.github.com> Date: Fri, 28 Jun 2024 18:37:51 -0300 Subject: [PATCH] Link Generic Thermostat helper entity to actuator entity device (#120767) --- .../components/generic_thermostat/__init__.py | 10 ++ .../components/generic_thermostat/climate.py | 11 ++- .../generic_thermostat/test_climate.py | 50 +++++++++- .../generic_thermostat/test_init.py | 98 +++++++++++++++++++ 4 files changed, 165 insertions(+), 4 deletions(-) create mode 100644 tests/components/generic_thermostat/test_init.py diff --git a/homeassistant/components/generic_thermostat/__init__.py b/homeassistant/components/generic_thermostat/__init__.py index 6a59e24ebd29..fcec36b8d35a 100644 --- a/homeassistant/components/generic_thermostat/__init__.py +++ b/homeassistant/components/generic_thermostat/__init__.py @@ -3,13 +3,23 @@ from homeassistant.config_entries import ConfigEntry from homeassistant.const import Platform from homeassistant.core import HomeAssistant +from homeassistant.helpers.device import ( + async_remove_stale_devices_links_keep_entity_device, +) +CONF_HEATER = "heater" DOMAIN = "generic_thermostat" PLATFORMS = [Platform.CLIMATE] async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up from a config entry.""" + + async_remove_stale_devices_links_keep_entity_device( + hass, + entry.entry_id, + entry.options[CONF_HEATER], + ) await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) entry.async_on_unload(entry.add_update_listener(config_entry_update_listener)) return True diff --git a/homeassistant/components/generic_thermostat/climate.py b/homeassistant/components/generic_thermostat/climate.py index d284c7d77725..1b19def9cf46 100644 --- a/homeassistant/components/generic_thermostat/climate.py +++ b/homeassistant/components/generic_thermostat/climate.py @@ -54,6 +54,7 @@ from homeassistant.core import ( ) from homeassistant.exceptions import ConditionError from homeassistant.helpers import condition, 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.event import ( async_track_state_change_event, @@ -63,14 +64,12 @@ from homeassistant.helpers.reload import async_setup_reload_service from homeassistant.helpers.restore_state import RestoreEntity from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType, VolDictType -from . import DOMAIN, PLATFORMS +from . import CONF_HEATER, DOMAIN, PLATFORMS _LOGGER = logging.getLogger(__name__) DEFAULT_TOLERANCE = 0.3 DEFAULT_NAME = "Generic Thermostat" - -CONF_HEATER = "heater" CONF_SENSOR = "target_sensor" CONF_MIN_TEMP = "min_temp" CONF_MAX_TEMP = "max_temp" @@ -190,6 +189,7 @@ async def _async_setup_config( async_add_entities( [ GenericThermostat( + hass, name, heater_entity_id, sensor_entity_id, @@ -220,6 +220,7 @@ class GenericThermostat(ClimateEntity, RestoreEntity): def __init__( self, + hass: HomeAssistant, name: str, heater_entity_id: str, sensor_entity_id: str, @@ -242,6 +243,10 @@ class GenericThermostat(ClimateEntity, RestoreEntity): self._attr_name = name self.heater_entity_id = heater_entity_id self.sensor_entity_id = sensor_entity_id + self._attr_device_info = async_device_info_to_link_from_entity( + hass, + heater_entity_id, + ) self.ac_mode = ac_mode self.min_cycle_duration = min_cycle_duration self._cold_tolerance = cold_tolerance diff --git a/tests/components/generic_thermostat/test_climate.py b/tests/components/generic_thermostat/test_climate.py index 1ecde733f483..7fb3e11e1890 100644 --- a/tests/components/generic_thermostat/test_climate.py +++ b/tests/components/generic_thermostat/test_climate.py @@ -44,12 +44,13 @@ from homeassistant.core import ( callback, ) from homeassistant.exceptions import ServiceValidationError -from homeassistant.helpers import entity_registry as er +from homeassistant.helpers import device_registry as dr, entity_registry as er from homeassistant.setup import async_setup_component from homeassistant.util import dt as dt_util from homeassistant.util.unit_system import METRIC_SYSTEM, US_CUSTOMARY_SYSTEM from tests.common import ( + MockConfigEntry, assert_setup_component, async_fire_time_changed, async_mock_service, @@ -1431,3 +1432,50 @@ async def test_reload(hass: HomeAssistant) -> None: assert len(hass.states.async_all()) == 1 assert hass.states.get("climate.test") is None assert hass.states.get("climate.reload") + + +async def test_device_id( + hass: HomeAssistant, + entity_registry: er.EntityRegistry, + device_registry: dr.DeviceRegistry, +) -> None: + """Test for source entity device.""" + + 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={("switch", "identifier_test")}, + connections={("mac", "30:31:32:33:34:35")}, + ) + source_entity = entity_registry.async_get_or_create( + "switch", + "test", + "source", + config_entry=source_config_entry, + device_id=source_device_entry.id, + ) + await hass.async_block_till_done() + assert entity_registry.async_get("switch.test_source") is not None + + helper_config_entry = MockConfigEntry( + data={}, + domain=GENERIC_THERMOSTAT_DOMAIN, + options={ + "name": "Test", + "heater": "switch.test_source", + "target_sensor": ENT_SENSOR, + "ac_mode": False, + "cold_tolerance": 0.3, + "hot_tolerance": 0.3, + }, + title="Test", + ) + helper_config_entry.add_to_hass(hass) + + assert await hass.config_entries.async_setup(helper_config_entry.entry_id) + await hass.async_block_till_done() + + helper_entity = entity_registry.async_get("climate.test") + assert helper_entity is not None + assert helper_entity.device_id == source_entity.device_id diff --git a/tests/components/generic_thermostat/test_init.py b/tests/components/generic_thermostat/test_init.py new file mode 100644 index 000000000000..0d6e106237ce --- /dev/null +++ b/tests/components/generic_thermostat/test_init.py @@ -0,0 +1,98 @@ +"""Test Generic Thermostat component setup process.""" + +from __future__ import annotations + +from homeassistant.components.generic_thermostat import DOMAIN +from homeassistant.core import HomeAssistant +from homeassistant.helpers import device_registry as dr, entity_registry as er + +from tests.common import MockConfigEntry + + +async def test_device_cleaning( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, +) -> None: + """Test cleaning of devices linked to the helper config entry.""" + + # 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={("switch", "identifier_test1")}, + connections={("mac", "30:31:32:33:34:01")}, + ) + + # Source entity registry + source_entity = entity_registry.async_get_or_create( + "switch", + "test", + "source", + config_entry=source_config_entry, + device_id=source_device1_entry.id, + ) + await hass.async_block_till_done() + assert entity_registry.async_get("switch.test_source") is not None + + # Configure the configuration entry for helper + helper_config_entry = MockConfigEntry( + data={}, + domain=DOMAIN, + options={ + "name": "Test", + "heater": "switch.test_source", + "target_sensor": "sensor.temperature", + "ac_mode": False, + "cold_tolerance": 0.3, + "hot_tolerance": 0.3, + }, + title="Test", + ) + helper_config_entry.add_to_hass(hass) + assert await hass.config_entries.async_setup(helper_config_entry.entry_id) + await hass.async_block_till_done() + + # Confirm the link between the source entity device and the helper entity + helper_entity = entity_registry.async_get("climate.test") + assert helper_entity is not None + assert helper_entity.device_id == source_entity.device_id + + # Device entry incorrectly linked to config entry + device_registry.async_get_or_create( + config_entry_id=helper_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=helper_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, 3 devices are expected to be linked + devices_before_reload = device_registry.devices.get_devices_for_config_entry_id( + helper_config_entry.entry_id + ) + assert len(devices_before_reload) == 3 + + # Config entry reload + await hass.config_entries.async_reload(helper_config_entry.entry_id) + await hass.async_block_till_done() + + # Confirm the link between the source entity device and the helper entity + helper_entity = entity_registry.async_get("climate.test") + assert helper_entity is not None + assert helper_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( + helper_config_entry.entry_id + ) + assert len(devices_after_reload) == 1 + + assert devices_after_reload[0].id == source_device1_entry.id