diff --git a/homeassistant/components/integration/__init__.py b/homeassistant/components/integration/__init__.py index effa0c4df55..4ccf0dec258 100644 --- a/homeassistant/components/integration/__init__.py +++ b/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) diff --git a/homeassistant/components/integration/sensor.py b/homeassistant/components/integration/sensor.py index d201fab0c6f..ffb7a3d8e6a 100644 --- a/homeassistant/components/integration/sensor.py +++ b/homeassistant/components/integration/sensor.py @@ -38,11 +38,8 @@ from homeassistant.core import ( State, callback, ) -from homeassistant.helpers import ( - config_validation as cv, - device_registry as dr, - entity_registry as er, -) +from homeassistant.helpers import config_validation as cv, entity_registry as er +from homeassistant.helpers.device import async_device_info_to_link_from_entity 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) - dev_reg = dr.async_get(hass) - # Resolve source entity device - 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 + device_info = async_device_info_to_link_from_entity( + hass, + source_entity_id, + ) if (unit_prefix := config_entry.options.get(CONF_UNIT_PREFIX)) == "none": # Before we had support for optional selectors, "none" was used for selecting nothing diff --git a/tests/components/integration/test_init.py b/tests/components/integration/test_init.py index 2ed32c7645c..9fee54f4500 100644 --- a/tests/components/integration/test_init.py +++ b/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