Fix entity subscribing to the old device after a device_id change (#183872)

homeassistant/helpers/entity.py

@@ -1649,9 +1649,6 @@ class Entity(
         if data["action"] != "update":
             return
 
-        if "device_id" in data["changes"]:
-            self._async_subscribe_device_updates()
-
         # Invalidate friendly name cache if relevant fields changed
         changes = data["changes"]
         if "name" in changes or "has_entity_name" in changes or "device_id" in changes:
@@ -1663,6 +1660,9 @@ class Entity(
         assert registry_entry is not None
         self.registry_entry = registry_entry
 
+        if "device_id" in changes:
+            self._async_subscribe_device_updates()
+
         if device_id := registry_entry.device_id:
             self.device_entry = dr.async_get(self.hass).async_get(device_id)
         else:

tests/helpers/test_entity.py

@@ -1708,6 +1708,113 @@ async def test_friendly_name_updated(
     assert state.attributes.get(ATTR_FRIENDLY_NAME) == expected_friendly_name3
 
 
+@pytest.mark.parametrize(
+    (
+        "entity_kwargs",
+        "new_device_identifier",
+        "expected_device_identifier",
+        "expected_friendly_name_after_move",
+        "expected_friendly_name_final",
+    ),
+    [
+        pytest.param(
+            {},
+            ("test", "new"),
+            ("test", "new"),
+            "New Device Entity",
+            "New renamed Entity",
+            id="no_device_to_device",
+        ),
+        pytest.param(
+            {"device_info": {"identifiers": {("test", "old")}}},
+            ("test", "new"),
+            ("test", "new"),
+            "New Device Entity",
+            "New renamed Entity",
+            id="device_to_other_device",
+        ),
+        pytest.param(
+            {"device_info": {"identifiers": {("test", "old")}}},
+            None,
+            None,
+            "Entity",
+            "Entity",
+            id="device_to_no_device",
+        ),
+    ],
+)
+async def test_device_updates_follow_device_id_change(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
+    entity_kwargs: dict[str, Any],
+    new_device_identifier: tuple[str, str] | None,
+    expected_device_identifier: tuple[str, str] | None,
+    expected_friendly_name_after_move: str,
+    expected_friendly_name_final: str,
+) -> None:
+    """Test device registry updates track the entity's current device.
+
+    After a device_id change, updates of the new device must be applied, and
+    updates of the previous device must not be.
+    """
+    config_entry = MockConfigEntry(entry_id="super-mock-id")
+    config_entry.add_to_hass(hass)
+    old_device = device_registry.async_get_or_create(
+        config_entry_id=config_entry.entry_id,
+        identifiers={("test", "old")},
+        name="Old Device",
+    )
+    new_device = device_registry.async_get_or_create(
+        config_entry_id=config_entry.entry_id,
+        identifiers={("test", "new")},
+        name="New Device",
+    )
+    devices = {("test", "old"): old_device, ("test", "new"): new_device, None: None}
+    device_ids = {("test", "old"): old_device.id, ("test", "new"): new_device.id}
+
+    ent = MockEntity(
+        unique_id="qwer", has_entity_name=True, name="Entity", **entity_kwargs
+    )
+
+    async def async_setup_entry(
+        hass: HomeAssistant,
+        config_entry: ConfigEntry,
+        async_add_entities: AddConfigEntryEntitiesCallback,
+    ) -> None:
+        """Mock setup entry method."""
+        async_add_entities([ent])
+
+    platform = MockPlatform(async_setup_entry=async_setup_entry)
+    entity_platform = MockEntityPlatform(
+        hass, platform_name=config_entry.domain, platform=platform
+    )
+    assert await entity_platform.async_setup_entry(config_entry)
+    await hass.async_block_till_done()
+
+    entity_registry.async_update_entity(
+        ent.entity_id, device_id=device_ids.get(new_device_identifier)
+    )
+    await hass.async_block_till_done()
+    expected_device = devices[expected_device_identifier]
+    assert ent.device_entry == expected_device
+    state = hass.states.get(ent.entity_id)
+    assert state.attributes[ATTR_FRIENDLY_NAME] == expected_friendly_name_after_move
+
+    # Renaming the previous device must not affect the entity
+    device_registry.async_update_device(old_device.id, name_by_user="Old renamed")
+    await hass.async_block_till_done()
+    assert ent.device_entry == expected_device
+    state = hass.states.get(ent.entity_id)
+    assert state.attributes[ATTR_FRIENDLY_NAME] == expected_friendly_name_after_move
+
+    # Renaming the new device must update the entity
+    device_registry.async_update_device(new_device.id, name_by_user="New renamed")
+    await hass.async_block_till_done()
+    state = hass.states.get(ent.entity_id)
+    assert state.attributes[ATTR_FRIENDLY_NAME] == expected_friendly_name_final
+
+
 async def test_device_entry_cleared_when_detached_from_device(
     hass: HomeAssistant,
     device_registry: dr.DeviceRegistry,