From cef7def024ea8876a36d88b69e9433ce5ceee93f Mon Sep 17 00:00:00 2001
From: dougiteixeira <31328123+dougiteixeira@users.noreply.github.com>
Date: Fri, 28 Jun 2024 18:38:45 -0300
Subject: [PATCH] Link Generic Hygrostat helper entity to actuator entity
device (#120759)
---
.../components/generic_hygrostat/__init__.py | 10 ++
.../generic_hygrostat/humidifier.py | 7 ++
.../generic_hygrostat/test_humidifier.py | 53 ++++++++-
.../components/generic_hygrostat/test_init.py | 102 ++++++++++++++++++
4 files changed, 171 insertions(+), 1 deletion(-)
create mode 100644 tests/components/generic_hygrostat/test_init.py
diff --git a/homeassistant/components/generic_hygrostat/__init__.py b/homeassistant/components/generic_hygrostat/__init__.py
index ef032da1ee2..b4a6014c5a4 100644
--- a/homeassistant/components/generic_hygrostat/__init__.py
+++ b/homeassistant/components/generic_hygrostat/__init__.py
@@ -7,6 +7,9 @@ from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_NAME, CONF_UNIQUE_ID, Platform
from homeassistant.core import HomeAssistant
from homeassistant.helpers import config_validation as cv, discovery
+from homeassistant.helpers.device import (
+ async_remove_stale_devices_links_keep_entity_device,
+)
from homeassistant.helpers.typing import ConfigType
DOMAIN = "generic_hygrostat"
@@ -78,6 +81,13 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
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_HUMIDIFIER],
+ )
+
await hass.config_entries.async_forward_entry_setups(entry, (Platform.HUMIDIFIER,))
entry.async_on_unload(entry.add_update_listener(config_entry_update_listener))
return True
diff --git a/homeassistant/components/generic_hygrostat/humidifier.py b/homeassistant/components/generic_hygrostat/humidifier.py
index a1f9936fa33..cc04dbf13c3 100644
--- a/homeassistant/components/generic_hygrostat/humidifier.py
+++ b/homeassistant/components/generic_hygrostat/humidifier.py
@@ -41,6 +41,7 @@ from homeassistant.core import (
callback,
)
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,
@@ -139,6 +140,7 @@ async def _async_setup_config(
async_add_entities(
[
GenericHygrostat(
+ hass,
name,
switch_entity_id,
sensor_entity_id,
@@ -167,6 +169,7 @@ class GenericHygrostat(HumidifierEntity, RestoreEntity):
def __init__(
self,
+ hass: HomeAssistant,
name: str,
switch_entity_id: str,
sensor_entity_id: str,
@@ -188,6 +191,10 @@ class GenericHygrostat(HumidifierEntity, RestoreEntity):
self._name = name
self._switch_entity_id = switch_entity_id
self._sensor_entity_id = sensor_entity_id
+ self._attr_device_info = async_device_info_to_link_from_entity(
+ hass,
+ switch_entity_id,
+ )
self._device_class = device_class or HumidifierDeviceClass.HUMIDIFIER
self._min_cycle_duration = min_cycle_duration
self._dry_tolerance = dry_tolerance
diff --git a/tests/components/generic_hygrostat/test_humidifier.py b/tests/components/generic_hygrostat/test_humidifier.py
index eadc1b22527..15d80885d27 100644
--- a/tests/components/generic_hygrostat/test_humidifier.py
+++ b/tests/components/generic_hygrostat/test_humidifier.py
@@ -7,6 +7,9 @@ import pytest
import voluptuous as vol
from homeassistant.components import input_boolean, switch
+from homeassistant.components.generic_hygrostat import (
+ DOMAIN as GENERIC_HYDROSTAT_DOMAIN,
+)
from homeassistant.components.humidifier import (
ATTR_HUMIDITY,
DOMAIN,
@@ -32,11 +35,12 @@ from homeassistant.core import (
State,
callback,
)
-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
import homeassistant.util.dt as dt_util
from tests.common import (
+ MockConfigEntry,
assert_setup_component,
async_fire_time_changed,
mock_restore_cache,
@@ -1782,3 +1786,50 @@ async def test_sensor_stale_duration(
# Not turning on by itself
assert hass.states.get(humidifier_switch).state == STATE_OFF
+
+
+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_HYDROSTAT_DOMAIN,
+ options={
+ "device_class": "humidifier",
+ "dry_tolerance": 2.0,
+ "humidifier": "switch.test_source",
+ "name": "Test",
+ "target_sensor": ENT_SENSOR,
+ "wet_tolerance": 4.0,
+ },
+ 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("humidifier.test")
+ assert helper_entity is not None
+ assert helper_entity.device_id == source_entity.device_id
diff --git a/tests/components/generic_hygrostat/test_init.py b/tests/components/generic_hygrostat/test_init.py
new file mode 100644
index 00000000000..bd4792f939d
--- /dev/null
+++ b/tests/components/generic_hygrostat/test_init.py
@@ -0,0 +1,102 @@
+"""Test Generic Hygrostat component setup process."""
+
+from __future__ import annotations
+
+from homeassistant.components.generic_hygrostat import (
+ DOMAIN as GENERIC_HYDROSTAT_DOMAIN,
+)
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import device_registry as dr, entity_registry as er
+
+from .test_humidifier import ENT_SENSOR
+
+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=GENERIC_HYDROSTAT_DOMAIN,
+ options={
+ "device_class": "humidifier",
+ "dry_tolerance": 2.0,
+ "humidifier": "switch.test_source",
+ "name": "Test",
+ "target_sensor": ENT_SENSOR,
+ "wet_tolerance": 4.0,
+ },
+ 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("humidifier.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("humidifier.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
--
GitLab