Link Generic Hygrostat helper entity to actuator entity device (#120759)

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

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

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

tests/components/generic_hygrostat/test_init.py

+"""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