Use the new device helpers in Integral (#120157)

homeassistant/components/integration/__init__.py

@@ -5,11 +5,22 @@ from __future__ import annotations
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import device_registry as dr
+from homeassistant.helpers.device import (
+    async_remove_stale_devices_links_keep_entity_device,
+)
+
+from .const import CONF_SOURCE_SENSOR
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Integration from a config entry."""
+
+    async_remove_stale_devices_links_keep_entity_device(
+        hass,
+        entry.entry_id,
+        entry.options[CONF_SOURCE_SENSOR],
+    )
+
     await hass.config_entries.async_forward_entry_setups(entry, (Platform.SENSOR,))
     entry.async_on_unload(entry.add_update_listener(config_entry_update_listener))
     return True
@@ -19,14 +30,6 @@ async def config_entry_update_listener(hass: HomeAssistant, entry: ConfigEntry)
     """Update listener, called when the config entry options are changed."""
     # Remove device link for entry, the source device may have changed.
     # The link will be recreated after load.
-    device_registry = dr.async_get(hass)
-    devices = device_registry.devices.get_devices_for_config_entry_id(entry.entry_id)
-
-    for device in devices:
-        device_registry.async_update_device(
-            device.id, remove_config_entry_id=entry.entry_id
-        )
-
     await hass.config_entries.async_reload(entry.entry_id)
 
 

homeassistant/components/integration/sensor.py

@@ -38,11 +38,8 @@ from homeassistant.core import (
     State,
     callback,
 )
-from homeassistant.helpers import (
+from homeassistant.helpers import config_validation as cv, entity_registry as er
-    config_validation as cv,
+from homeassistant.helpers.device import async_device_info_to_link_from_entity
-    device_registry as dr,
-    entity_registry as er,
-)
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.event import async_call_later, async_track_state_change_event
@@ -249,27 +246,10 @@ async def async_setup_entry(
         registry, config_entry.options[CONF_SOURCE_SENSOR]
     )
 
-    source_entity = er.EntityRegistry.async_get(registry, source_entity_id)
+    device_info = async_device_info_to_link_from_entity(
-    dev_reg = dr.async_get(hass)
+        hass,
-    # Resolve source entity device
+        source_entity_id,
-    if (
+    )
-        (source_entity is not None)
-        and (source_entity.device_id is not None)
-        and (
-            (
-                device := dev_reg.async_get(
-                    device_id=source_entity.device_id,
-                )
-            )
-            is not None
-        )
-    ):
-        device_info = DeviceInfo(
-            identifiers=device.identifiers,
-            connections=device.connections,
-        )
-    else:
-        device_info = None
 
     if (unit_prefix := config_entry.options.get(CONF_UNIT_PREFIX)) == "none":
         # Before we had support for optional selectors, "none" was used for selecting nothing

tests/components/integration/test_init.py

@@ -121,3 +121,91 @@ async def test_entry_changed(hass: HomeAssistant, platform) -> None:
     # Check that the config entry association has updated
     assert config_entry.entry_id not in _get_device_config_entries(input_entry)
     assert config_entry.entry_id in _get_device_config_entries(valid_entry)
+
+
+async def test_device_cleaning(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
+) -> None:
+    """Test for source entity device for Integration."""
+
+    # 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 Integration
+    integration_config_entry = MockConfigEntry(
+        data={},
+        domain=DOMAIN,
+        options={
+            "method": "trapezoidal",
+            "name": "Integration",
+            "round": 1.0,
+            "source": "sensor.test_source",
+            "unit_prefix": "k",
+            "unit_time": "min",
+            "max_sub_interval": {"minutes": 1},
+        },
+        title="Integration",
+    )
+    integration_config_entry.add_to_hass(hass)
+    assert await hass.config_entries.async_setup(integration_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    # Confirm the link between the source entity device and the integration sensor
+    integration_entity = entity_registry.async_get("sensor.integration")
+    assert integration_entity is not None
+    assert integration_entity.device_id == source_entity.device_id
+
+    # Device entry incorrectly linked to Integration config entry
+    device_registry.async_get_or_create(
+        config_entry_id=integration_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=integration_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(
+        integration_config_entry.entry_id
+    )
+    assert len(devices_before_reload) == 3
+
+    # Config entry reload
+    await hass.config_entries.async_reload(integration_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    # Confirm the link between the source entity device and the integration sensor after reload
+    integration_entity = entity_registry.async_get("sensor.integration")
+    assert integration_entity is not None
+    assert integration_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(
+        integration_config_entry.entry_id
+    )
+    assert len(devices_after_reload) == 1