Make async_remove_stale_devices_links_keep_entity_device move entities (#145719)

Co-authored-by: Jan Bouwhuis <jbouwh@users.noreply.github.com>
Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>

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
+    Also moves all entities linked to the entry_id to the device of
-    linked to the configuration entry will be maintained.
+    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.
+    """Remove entry_id from all devices except current_device_id."""
-
-    Only the device passed in the current_device_id parameter linked to
-    the configuration entry will be maintained.
-    """
 
     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:

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")
+    helper_config_entry = MockConfigEntry(domain="helper_integration")
-    config_entry.add_to_hass(hass)
+    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(source_entity.entity_id) is not None
-    assert entity_registry.async_get("sensor.test_source") is not None
 
-    devices_config_entry = device_registry.devices.get_devices_for_config_entry_id(
+    # Helper entity connected to a stale device
-        config_entry.entry_id
+    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(
+    await hass.async_block_till_done()
-        config_entry.entry_id
+
+    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
+    # After cleanup, only one device is expected to be linked to the config entry, and
-    assert len(devices_config_entry) == 1
+    # the entities should exist and be linked to the current device
-
+    assert len(devices_helper_entry) == 1
-    assert current_device in devices_config_entry
+    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(