From 9d4375ca769b73f090c594e36844d733a9d68041 Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Tue, 27 May 2025 23:00:52 +0200 Subject: [PATCH] Make async_remove_stale_devices_links_keep_entity_device move entities (#145719) Co-authored-by: Jan Bouwhuis Co-authored-by: Joost Lekkerkerker --- homeassistant/helpers/device.py | 20 ++++++----- tests/helpers/test_device.py | 60 ++++++++++++++++++++------------- 2 files changed, 49 insertions(+), 31 deletions(-) diff --git a/homeassistant/helpers/device.py b/homeassistant/helpers/device.py index 16212422236..a7d888900b1 100644 --- a/homeassistant/helpers/device.py +++ b/homeassistant/helpers/device.py @@ -64,10 +64,10 @@ def async_remove_stale_devices_links_keep_entity_device( entry_id: str, source_entity_id_or_uuid: str, ) -> None: - """Remove the link between stale devices and a configuration entry. + """Remove entry_id from all devices except that of source_entity_id_or_uuid. - Only the device passed in the source_entity_id_or_uuid parameter - linked to the configuration entry will be maintained. + Also moves all entities linked to the entry_id to the device of + source_entity_id_or_uuid. """ async_remove_stale_devices_links_keep_current_device( @@ -83,13 +83,17 @@ def async_remove_stale_devices_links_keep_current_device( entry_id: str, current_device_id: str | None, ) -> None: - """Remove the link between stale devices and a configuration entry. - - Only the device passed in the current_device_id parameter linked to - the configuration entry will be maintained. - """ + """Remove entry_id from all devices except current_device_id.""" dev_reg = dr.async_get(hass) + ent_reg = er.async_get(hass) + + # Make sure all entities are linked to the correct device + for entity in ent_reg.entities.get_entries_for_config_entry_id(entry_id): + if entity.device_id == current_device_id: + continue + ent_reg.async_update_entity(entity.entity_id, device_id=current_device_id) + # Removes all devices from the config entry that are not the same as the current device for device in dev_reg.devices.get_devices_for_config_entry_id(entry_id): if device.id == current_device_id: diff --git a/tests/helpers/test_device.py b/tests/helpers/test_device.py index 852d418da23..266435ef05d 100644 --- a/tests/helpers/test_device.py +++ b/tests/helpers/test_device.py @@ -118,61 +118,75 @@ async def test_remove_stale_device_links_keep_entity_device( entity_registry: er.EntityRegistry, ) -> None: """Test cleaning works for entity.""" - config_entry = MockConfigEntry(domain="hue") - config_entry.add_to_hass(hass) + helper_config_entry = MockConfigEntry(domain="helper_integration") + helper_config_entry.add_to_hass(hass) + host_config_entry = MockConfigEntry(domain="host_integration") + host_config_entry.add_to_hass(hass) current_device = device_registry.async_get_or_create( identifiers={("test", "current_device")}, connections={("mac", "30:31:32:33:34:00")}, - config_entry_id=config_entry.entry_id, + config_entry_id=helper_config_entry.entry_id, ) - assert current_device is not None - device_registry.async_get_or_create( + stale_device_1 = device_registry.async_get_or_create( identifiers={("test", "stale_device_1")}, connections={("mac", "30:31:32:33:34:01")}, - config_entry_id=config_entry.entry_id, + config_entry_id=helper_config_entry.entry_id, ) device_registry.async_get_or_create( identifiers={("test", "stale_device_2")}, connections={("mac", "30:31:32:33:34:02")}, - config_entry_id=config_entry.entry_id, + config_entry_id=helper_config_entry.entry_id, ) - # Source entity registry + # Source entity source_entity = entity_registry.async_get_or_create( "sensor", - "test", + "host_integration", "source", - config_entry=config_entry, + config_entry=host_config_entry, device_id=current_device.id, ) - await hass.async_block_till_done() - assert entity_registry.async_get("sensor.test_source") is not None + assert entity_registry.async_get(source_entity.entity_id) is not None - devices_config_entry = device_registry.devices.get_devices_for_config_entry_id( - config_entry.entry_id + # Helper entity connected to a stale device + helper_entity = entity_registry.async_get_or_create( + "sensor", + "helper_integration", + "helper", + config_entry=helper_config_entry, + device_id=stale_device_1.id, + ) + assert entity_registry.async_get(helper_entity.entity_id) is not None + + devices_helper_entry = device_registry.devices.get_devices_for_config_entry_id( + helper_config_entry.entry_id ) # 3 devices linked to the config entry are expected (1 current device + 2 stales) - assert len(devices_config_entry) == 3 + assert len(devices_helper_entry) == 3 - # Manual cleanup should unlink stales devices from the config entry + # Manual cleanup should unlink stale devices from the config entry async_remove_stale_devices_links_keep_entity_device( hass, - entry_id=config_entry.entry_id, + entry_id=helper_config_entry.entry_id, source_entity_id_or_uuid=source_entity.entity_id, ) - devices_config_entry = device_registry.devices.get_devices_for_config_entry_id( - config_entry.entry_id + await hass.async_block_till_done() + + devices_helper_entry = device_registry.devices.get_devices_for_config_entry_id( + helper_config_entry.entry_id ) - # After cleanup, only one device is expected to be linked to the config entry - assert len(devices_config_entry) == 1 - - assert current_device in devices_config_entry + # After cleanup, only one device is expected to be linked to the config entry, and + # the entities should exist and be linked to the current device + assert len(devices_helper_entry) == 1 + assert current_device in devices_helper_entry + assert entity_registry.async_get(source_entity.entity_id) is not None + assert entity_registry.async_get(helper_entity.entity_id) is not None async def test_remove_stale_devices_links_keep_current_device(