Handle changes to source entity in derivative helper (#146407)

* Handle changes to source entity in derivative helper

* Rename helper function, improve docstring

* Add tests

* Improve derivative tests

* Deduplicate tests

* Rename helpers/helper_entity.py to helpers/helper_integration.py

* Rename tests

homeassistant/components/derivative/__init__.py

@@ -2,12 +2,18 @@
 
 from __future__ import annotations
 
+from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_SOURCE, Platform
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.device import (
+    async_entity_id_to_device_id,
     async_remove_stale_devices_links_keep_entity_device,
 )
+from homeassistant.helpers.helper_integration import async_handle_source_entity_changes
+
+from .const import DOMAIN
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
@@ -17,6 +23,27 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         hass, entry.entry_id, entry.options[CONF_SOURCE]
     )
 
+    def set_source_entity_id_or_uuid(source_entity_id: str) -> None:
+        hass.config_entries.async_update_entry(
+            entry,
+            options={**entry.options, CONF_SOURCE: source_entity_id},
+        )
+
+    entity_registry = er.async_get(hass)
+    entry.async_on_unload(
+        async_handle_source_entity_changes(
+            hass,
+            helper_config_entry_id=entry.entry_id,
+            get_helper_entity_id=lambda: entity_registry.async_get_entity_id(
+                SENSOR_DOMAIN, DOMAIN, entry.entry_id
+            ),
+            set_source_entity_id_or_uuid=set_source_entity_id_or_uuid,
+            source_device_id=async_entity_id_to_device_id(
+                hass, entry.options[CONF_SOURCE]
+            ),
+            source_entity_id_or_uuid=entry.options[CONF_SOURCE],
+        )
+    )
     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

homeassistant/components/switch_as_x/__init__.py

@@ -9,9 +9,9 @@ import voluptuous as vol
 from homeassistant.components.homeassistant import exposed_entities
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_ENTITY_ID
-from homeassistant.core import Event, HomeAssistant, callback, valid_entity_id
+from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import device_registry as dr, entity_registry as er
-from homeassistant.helpers.event import async_track_entity_registry_updated_event
+from homeassistant.helpers.helper_integration import async_handle_source_entity_changes
 
 from .const import CONF_INVERT, CONF_TARGET_DOMAIN, DOMAIN
 
@@ -42,7 +42,6 @@ def async_add_to_device(
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up a config entry."""
     entity_registry = er.async_get(hass)
-    device_registry = dr.async_get(hass)
     try:
         entity_id = er.async_validate_entity_id(
             entity_registry, entry.options[CONF_ENTITY_ID]
@@ -55,72 +54,26 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         )
         return False
 
-    async def async_registry_updated(
+    def set_source_entity_id_or_uuid(source_entity_id: str) -> None:
-        event: Event[er.EventEntityRegistryUpdatedData],
+        hass.config_entries.async_update_entry(
-    ) -> None:
+            entry,
-        """Handle entity registry update."""
+            options={**entry.options, CONF_ENTITY_ID: source_entity_id},
-        data = event.data
+        )
-        if data["action"] == "remove":
-            await hass.config_entries.async_remove(entry.entry_id)
-
-        if data["action"] != "update":
-            return
-
-        if "entity_id" in data["changes"]:
-            # Entity_id changed, update or reload the config entry
-            if valid_entity_id(entry.options[CONF_ENTITY_ID]):
-                # If the entity is pointed to by an entity ID, update the entry
-                hass.config_entries.async_update_entry(
-                    entry,
-                    options={**entry.options, CONF_ENTITY_ID: data["entity_id"]},
-                )
-            else:
-                await hass.config_entries.async_reload(entry.entry_id)
-
-        if device_id and "device_id" in data["changes"]:
-            # Handle the wrapped switch being moved to a different device or removed
-            # from the device
-            if (
-                not (entity_entry := entity_registry.async_get(data[CONF_ENTITY_ID]))
-                or not device_registry.async_get(device_id)
-                or entity_entry.device_id == device_id
-            ):
-                # No need to do any cleanup
-                return
-
-            # The wrapped switch has been moved to a different device, update the
-            # switch_as_x entity and the device entry to include our config entry
-            switch_as_x_entity_id = entity_registry.async_get_entity_id(
-                entry.options[CONF_TARGET_DOMAIN], DOMAIN, entry.entry_id
-            )
-            if switch_as_x_entity_id:
-                # Update the switch_as_x entity to point to the new device (or no device)
-                entity_registry.async_update_entity(
-                    switch_as_x_entity_id, device_id=entity_entry.device_id
-                )
-
-            if entity_entry.device_id is not None:
-                device_registry.async_update_device(
-                    entity_entry.device_id, add_config_entry_id=entry.entry_id
-                )
-
-            device_registry.async_update_device(
-                device_id, remove_config_entry_id=entry.entry_id
-            )
-
-            # Reload the config entry so the switch_as_x entity is recreated with
-            # correct device info
-            await hass.config_entries.async_reload(entry.entry_id)
 
     entry.async_on_unload(
-        async_track_entity_registry_updated_event(
+        async_handle_source_entity_changes(
-            hass, entity_id, async_registry_updated
+            hass,
+            helper_config_entry_id=entry.entry_id,
+            get_helper_entity_id=lambda: entity_registry.async_get_entity_id(
+                entry.options[CONF_TARGET_DOMAIN], DOMAIN, entry.entry_id
+            ),
+            set_source_entity_id_or_uuid=set_source_entity_id_or_uuid,
+            source_device_id=async_add_to_device(hass, entry, entity_id),
+            source_entity_id_or_uuid=entry.options[CONF_ENTITY_ID],
         )
     )
     entry.async_on_unload(entry.add_update_listener(config_entry_update_listener))
 
-    device_id = async_add_to_device(hass, entry, entity_id)
-
     await hass.config_entries.async_forward_entry_setups(
         entry, (entry.options[CONF_TARGET_DOMAIN],)
     )

homeassistant/helpers/helper_integration.py

@@ -0,0 +1,105 @@
+"""Helpers for helper integrations."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+from homeassistant.core import CALLBACK_TYPE, Event, HomeAssistant, valid_entity_id
+
+from . import device_registry as dr, entity_registry as er
+from .event import async_track_entity_registry_updated_event
+
+
+def async_handle_source_entity_changes(
+    hass: HomeAssistant,
+    *,
+    helper_config_entry_id: str,
+    get_helper_entity_id: Callable[[], str | None],
+    set_source_entity_id_or_uuid: Callable[[str], None],
+    source_device_id: str | None,
+    source_entity_id_or_uuid: str,
+) -> CALLBACK_TYPE:
+    """Handle changes to a helper entity's source entity.
+
+    The following changes are handled:
+    - Entity removal: If the source entity is removed, the helper config entry
+      is removed, and the helper entity is cleaned up.
+    - Entity ID changed: If the source entity's entity ID changes and the source
+      entity is identified by an entity ID, the set_source_entity_id_or_uuid is
+      called. If the source entity is identified by a UUID, the helper config entry
+      is reloaded.
+    - Source entity moved to another device: The helper entity is updated to link
+      to the new device, and the helper config entry removed from the old device
+      and added to the new device. Then the helper config entry is reloaded.
+    - Source entity removed from the device: The helper entity is updated to link
+      to no device, and the helper config entry removed from the old device. Then
+      the helper config entry is reloaded.
+    """
+
+    async def async_registry_updated(
+        event: Event[er.EventEntityRegistryUpdatedData],
+    ) -> None:
+        """Handle entity registry update."""
+        nonlocal source_device_id
+
+        data = event.data
+        if data["action"] == "remove":
+            await hass.config_entries.async_remove(helper_config_entry_id)
+
+        if data["action"] != "update":
+            return
+
+        if "entity_id" in data["changes"]:
+            # Entity_id changed, update or reload the config entry
+            if valid_entity_id(source_entity_id_or_uuid):
+                # If the entity is pointed to by an entity ID, update the entry
+                set_source_entity_id_or_uuid(data["entity_id"])
+            else:
+                await hass.config_entries.async_reload(helper_config_entry_id)
+
+        if not source_device_id or "device_id" not in data["changes"]:
+            return
+
+        # Handle the source entity being moved to a different device or removed
+        # from the device
+        if (
+            not (source_entity_entry := entity_registry.async_get(data["entity_id"]))
+            or not device_registry.async_get(source_device_id)
+            or source_entity_entry.device_id == source_device_id
+        ):
+            # No need to do any cleanup
+            return
+
+        # The source entity has been moved to a different device, update the helper
+        # helper entity to link to the new device and the helper device to include
+        # the helper config entry
+        helper_entity_id = get_helper_entity_id()
+        if helper_entity_id:
+            # Update the helper entity to link to the new device (or no device)
+            entity_registry.async_update_entity(
+                helper_entity_id, device_id=source_entity_entry.device_id
+            )
+
+        if source_entity_entry.device_id is not None:
+            device_registry.async_update_device(
+                source_entity_entry.device_id,
+                add_config_entry_id=helper_config_entry_id,
+            )
+
+        device_registry.async_update_device(
+            source_device_id, remove_config_entry_id=helper_config_entry_id
+        )
+        source_device_id = source_entity_entry.device_id
+
+        # Reload the config entry so the helper entity is recreated with
+        # correct device info
+        await hass.config_entries.async_reload(helper_config_entry_id)
+
+    device_registry = dr.async_get(hass)
+    entity_registry = er.async_get(hass)
+    source_entity_id = er.async_validate_entity_id(
+        entity_registry, source_entity_id_or_uuid
+    )
+    return async_track_entity_registry_updated_event(
+        hass, source_entity_id, async_registry_updated
+    )

tests/components/derivative/test_init.py

@@ -1,23 +1,103 @@
 """Test the Derivative integration."""
 
+from unittest.mock import patch
+
 import pytest
 
+from homeassistant.components import derivative
+from homeassistant.components.derivative.config_flow import ConfigFlowHandler
 from homeassistant.components.derivative.const import DOMAIN
-from homeassistant.core import HomeAssistant
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import Event, HomeAssistant
 from homeassistant.helpers import device_registry as dr, entity_registry as er
+from homeassistant.helpers.event import async_track_entity_registry_updated_event
 
 from tests.common import MockConfigEntry
 
 
-@pytest.mark.parametrize("platform", ["sensor"])
+@pytest.fixture
+def sensor_config_entry(hass: HomeAssistant) -> er.RegistryEntry:
+    """Fixture to create a sensor config entry."""
+    sensor_config_entry = MockConfigEntry()
+    sensor_config_entry.add_to_hass(hass)
+    return sensor_config_entry
+
+
+@pytest.fixture
+def sensor_device(
+    device_registry: dr.DeviceRegistry, sensor_config_entry: ConfigEntry
+) -> dr.DeviceEntry:
+    """Fixture to create a sensor device."""
+    return device_registry.async_get_or_create(
+        config_entry_id=sensor_config_entry.entry_id,
+        connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+    )
+
+
+@pytest.fixture
+def sensor_entity_entry(
+    entity_registry: er.EntityRegistry,
+    sensor_config_entry: ConfigEntry,
+    sensor_device: dr.DeviceEntry,
+) -> er.RegistryEntry:
+    """Fixture to create a sensor entity entry."""
+    return entity_registry.async_get_or_create(
+        "sensor",
+        "test",
+        "unique",
+        config_entry=sensor_config_entry,
+        device_id=sensor_device.id,
+        original_name="ABC",
+    )
+
+
+@pytest.fixture
+def derivative_config_entry(
+    hass: HomeAssistant,
+    sensor_entity_entry: er.RegistryEntry,
+) -> MockConfigEntry:
+    """Fixture to create a derivative config entry."""
+    config_entry = MockConfigEntry(
+        data={},
+        domain=DOMAIN,
+        options={
+            "name": "My derivative",
+            "round": 1.0,
+            "source": sensor_entity_entry.entity_id,
+            "time_window": {"seconds": 0.0},
+            "unit_prefix": "k",
+            "unit_time": "min",
+        },
+        title="My derivative",
+        version=ConfigFlowHandler.VERSION,
+        minor_version=ConfigFlowHandler.MINOR_VERSION,
+    )
+
+    config_entry.add_to_hass(hass)
+
+    return config_entry
+
+
+def track_entity_registry_actions(hass: HomeAssistant, entity_id: str) -> list[str]:
+    """Track entity registry actions for an entity."""
+    events = []
+
+    def add_event(event: Event[er.EventEntityRegistryUpdatedData]) -> None:
+        """Add entity registry updated event to the list."""
+        events.append(event.data["action"])
+
+    async_track_entity_registry_updated_event(hass, entity_id, add_event)
+
+    return events
+
+
 async def test_setup_and_remove_config_entry(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
-    platform: str,
 ) -> None:
     """Test setting up and removing a config entry."""
     input_sensor_entity_id = "sensor.input"
-    derivative_entity_id = f"{platform}.my_derivative"
+    derivative_entity_id = "sensor.my_derivative"
 
     # Setup the config entry
     config_entry = MockConfigEntry(
@@ -147,3 +227,194 @@ async def test_device_cleaning(
         derivative_config_entry.entry_id
     )
     assert len(devices_after_reload) == 1
+
+
+async def test_async_handle_source_entity_changes_source_entity_removed(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
+    derivative_config_entry: MockConfigEntry,
+    sensor_config_entry: ConfigEntry,
+    sensor_device: dr.DeviceEntry,
+    sensor_entity_entry: er.RegistryEntry,
+) -> None:
+    """Test the derivative config entry is removed when the source entity is removed."""
+    # Add another config entry to the sensor device
+    other_config_entry = MockConfigEntry()
+    other_config_entry.add_to_hass(hass)
+    device_registry.async_update_device(
+        sensor_device.id, add_config_entry_id=other_config_entry.entry_id
+    )
+
+    assert await hass.config_entries.async_setup(derivative_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    derivative_entity_entry = entity_registry.async_get("sensor.my_derivative")
+    assert derivative_entity_entry.device_id == sensor_entity_entry.device_id
+
+    sensor_device = device_registry.async_get(sensor_device.id)
+    assert derivative_config_entry.entry_id in sensor_device.config_entries
+
+    events = track_entity_registry_actions(hass, derivative_entity_entry.entity_id)
+
+    # Remove the source sensor's config entry from the device, this removes the
+    # source sensor
+    with patch(
+        "homeassistant.components.derivative.async_unload_entry",
+        wraps=derivative.async_unload_entry,
+    ) as mock_unload_entry:
+        device_registry.async_update_device(
+            sensor_device.id, remove_config_entry_id=sensor_config_entry.entry_id
+        )
+        await hass.async_block_till_done()
+        await hass.async_block_till_done()
+    mock_unload_entry.assert_called_once()
+
+    # Check that the derivative config entry is removed from the device
+    sensor_device = device_registry.async_get(sensor_device.id)
+    assert derivative_config_entry.entry_id not in sensor_device.config_entries
+
+    # Check that the derivative config entry is removed
+    assert derivative_config_entry.entry_id not in hass.config_entries.async_entry_ids()
+
+    # Check we got the expected events
+    assert events == ["remove"]
+
+
+async def test_async_handle_source_entity_changes_source_entity_removed_from_device(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
+    derivative_config_entry: MockConfigEntry,
+    sensor_device: dr.DeviceEntry,
+    sensor_entity_entry: er.RegistryEntry,
+) -> None:
+    """Test the source entity removed from the source device."""
+    assert await hass.config_entries.async_setup(derivative_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    derivative_entity_entry = entity_registry.async_get("sensor.my_derivative")
+    assert derivative_entity_entry.device_id == sensor_entity_entry.device_id
+
+    sensor_device = device_registry.async_get(sensor_device.id)
+    assert derivative_config_entry.entry_id in sensor_device.config_entries
+
+    events = track_entity_registry_actions(hass, derivative_entity_entry.entity_id)
+
+    # Remove the source sensor from the device
+    with patch(
+        "homeassistant.components.derivative.async_unload_entry",
+        wraps=derivative.async_unload_entry,
+    ) as mock_unload_entry:
+        entity_registry.async_update_entity(
+            sensor_entity_entry.entity_id, device_id=None
+        )
+        await hass.async_block_till_done()
+    mock_unload_entry.assert_called_once()
+
+    # Check that the derivative config entry is removed from the device
+    sensor_device = device_registry.async_get(sensor_device.id)
+    assert derivative_config_entry.entry_id not in sensor_device.config_entries
+
+    # Check that the derivative config entry is not removed
+    assert derivative_config_entry.entry_id in hass.config_entries.async_entry_ids()
+
+    # Check we got the expected events
+    assert events == ["update"]
+
+
+async def test_async_handle_source_entity_changes_source_entity_moved_other_device(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
+    derivative_config_entry: MockConfigEntry,
+    sensor_config_entry: ConfigEntry,
+    sensor_device: dr.DeviceEntry,
+    sensor_entity_entry: er.RegistryEntry,
+) -> None:
+    """Test the source entity is moved to another device."""
+    sensor_device_2 = device_registry.async_get_or_create(
+        config_entry_id=sensor_config_entry.entry_id,
+        connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:FF")},
+    )
+
+    assert await hass.config_entries.async_setup(derivative_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    derivative_entity_entry = entity_registry.async_get("sensor.my_derivative")
+    assert derivative_entity_entry.device_id == sensor_entity_entry.device_id
+
+    sensor_device = device_registry.async_get(sensor_device.id)
+    assert derivative_config_entry.entry_id in sensor_device.config_entries
+    sensor_device_2 = device_registry.async_get(sensor_device_2.id)
+    assert derivative_config_entry.entry_id not in sensor_device_2.config_entries
+
+    events = track_entity_registry_actions(hass, derivative_entity_entry.entity_id)
+
+    # Move the source sensor to another device
+    with patch(
+        "homeassistant.components.derivative.async_unload_entry",
+        wraps=derivative.async_unload_entry,
+    ) as mock_unload_entry:
+        entity_registry.async_update_entity(
+            sensor_entity_entry.entity_id, device_id=sensor_device_2.id
+        )
+        await hass.async_block_till_done()
+    mock_unload_entry.assert_called_once()
+
+    # Check that the derivative config entry is moved to the other device
+    sensor_device = device_registry.async_get(sensor_device.id)
+    assert derivative_config_entry.entry_id not in sensor_device.config_entries
+    sensor_device_2 = device_registry.async_get(sensor_device_2.id)
+    assert derivative_config_entry.entry_id in sensor_device_2.config_entries
+
+    # Check that the derivative config entry is not removed
+    assert derivative_config_entry.entry_id in hass.config_entries.async_entry_ids()
+
+    # Check we got the expected events
+    assert events == ["update"]
+
+
+async def test_async_handle_source_entity_new_entity_id(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
+    derivative_config_entry: MockConfigEntry,
+    sensor_device: dr.DeviceEntry,
+    sensor_entity_entry: er.RegistryEntry,
+) -> None:
+    """Test the source entity's entity ID is changed."""
+    assert await hass.config_entries.async_setup(derivative_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    derivative_entity_entry = entity_registry.async_get("sensor.my_derivative")
+    assert derivative_entity_entry.device_id == sensor_entity_entry.device_id
+
+    sensor_device = device_registry.async_get(sensor_device.id)
+    assert derivative_config_entry.entry_id in sensor_device.config_entries
+
+    events = track_entity_registry_actions(hass, derivative_entity_entry.entity_id)
+
+    # Change the source entity's entity ID
+    with patch(
+        "homeassistant.components.derivative.async_unload_entry",
+        wraps=derivative.async_unload_entry,
+    ) as mock_unload_entry:
+        entity_registry.async_update_entity(
+            sensor_entity_entry.entity_id, new_entity_id="sensor.new_entity_id"
+        )
+        await hass.async_block_till_done()
+    mock_unload_entry.assert_called_once()
+
+    # Check that the derivative config entry is updated with the new entity ID
+    assert derivative_config_entry.options["source"] == "sensor.new_entity_id"
+
+    # Check that the helper config is still in the device
+    sensor_device = device_registry.async_get(sensor_device.id)
+    assert derivative_config_entry.entry_id in sensor_device.config_entries
+
+    # Check that the derivative config entry is not removed
+    assert derivative_config_entry.entry_id in hass.config_entries.async_entry_ids()
+
+    # Check we got the expected events
+    assert events == []

tests/helpers/test_helper_integration.py

@@ -0,0 +1,424 @@
+"""Tests for the helper entity helpers."""
+
+from collections.abc import Generator
+from unittest.mock import AsyncMock, Mock
+
+import pytest
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import Event, HomeAssistant
+from homeassistant.helpers import device_registry as dr, entity_registry as er
+from homeassistant.helpers.event import async_track_entity_registry_updated_event
+from homeassistant.helpers.helper_integration import async_handle_source_entity_changes
+
+from tests.common import (
+    MockConfigEntry,
+    MockModule,
+    mock_config_flow,
+    mock_integration,
+    mock_platform,
+)
+
+HELPER_DOMAIN = "helper"
+SOURCE_DOMAIN = "test"
+
+
+@pytest.fixture
+def source_config_entry(hass: HomeAssistant) -> er.RegistryEntry:
+    """Fixture to create a source config entry."""
+    source_config_entry = MockConfigEntry()
+    source_config_entry.add_to_hass(hass)
+    return source_config_entry
+
+
+@pytest.fixture
+def source_device(
+    device_registry: dr.DeviceRegistry,
+    source_config_entry: ConfigEntry,
+) -> dr.DeviceEntry:
+    """Fixture to create a source device."""
+    return device_registry.async_get_or_create(
+        config_entry_id=source_config_entry.entry_id,
+        connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+    )
+
+
+@pytest.fixture
+def source_entity_entry(
+    entity_registry: er.EntityRegistry,
+    source_config_entry: ConfigEntry,
+    source_device: dr.DeviceEntry,
+) -> er.RegistryEntry:
+    """Fixture to create a source entity entry."""
+    return entity_registry.async_get_or_create(
+        "sensor",
+        SOURCE_DOMAIN,
+        "unique",
+        config_entry=source_config_entry,
+        device_id=source_device.id,
+        original_name="ABC",
+    )
+
+
+@pytest.fixture
+def helper_config_entry(
+    hass: HomeAssistant,
+    source_entity_entry: er.RegistryEntry,
+    use_entity_registry_id: bool,
+) -> MockConfigEntry:
+    """Fixture to create a helper config entry."""
+    config_entry = MockConfigEntry(
+        data={},
+        domain=HELPER_DOMAIN,
+        options={
+            "name": "My helper",
+            "round": 1.0,
+            "source": source_entity_entry.id
+            if use_entity_registry_id
+            else source_entity_entry.entity_id,
+            "time_window": {"seconds": 0.0},
+            "unit_prefix": "k",
+            "unit_time": "min",
+        },
+        title="My helper",
+    )
+
+    config_entry.add_to_hass(hass)
+
+    return config_entry
+
+
+@pytest.fixture
+def mock_helper_flow() -> Generator[None]:
+    """Mock helper config flow."""
+
+    class MockConfigFlow:
+        """Mock the helper config flow."""
+
+        VERSION = 1
+        MINOR_VERSION = 1
+
+    with mock_config_flow(HELPER_DOMAIN, MockConfigFlow):
+        yield
+
+
+@pytest.fixture
+def helper_entity_entry(
+    entity_registry: er.EntityRegistry,
+    helper_config_entry: ConfigEntry,
+    source_device: dr.DeviceEntry,
+) -> er.RegistryEntry:
+    """Fixture to create a helper entity entry."""
+    return entity_registry.async_get_or_create(
+        "sensor",
+        HELPER_DOMAIN,
+        helper_config_entry.entry_id,
+        config_entry=helper_config_entry,
+        device_id=source_device.id,
+        original_name="ABC",
+    )
+
+
+@pytest.fixture
+def async_remove_entry() -> AsyncMock:
+    """Fixture to mock async_remove_entry."""
+    return AsyncMock(return_value=True)
+
+
+@pytest.fixture
+def async_unload_entry() -> AsyncMock:
+    """Fixture to mock async_unload_entry."""
+    return AsyncMock(return_value=True)
+
+
+@pytest.fixture
+def set_source_entity_id_or_uuid() -> AsyncMock:
+    """Fixture to mock async_unload_entry."""
+    return Mock()
+
+
+@pytest.fixture
+def mock_helper_integration(
+    hass: HomeAssistant,
+    entity_registry: er.EntityRegistry,
+    helper_config_entry: MockConfigEntry,
+    source_entity_entry: er.RegistryEntry,
+    async_remove_entry: AsyncMock,
+    async_unload_entry: AsyncMock,
+    set_source_entity_id_or_uuid: Mock,
+) -> None:
+    """Mock the helper integration."""
+
+    def get_helper_entity_id() -> str | None:
+        """Get the helper entity ID."""
+        return entity_registry.async_get_entity_id(
+            "sensor", HELPER_DOMAIN, helper_config_entry.entry_id
+        )
+
+    async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+        """Mock setup entry."""
+        async_handle_source_entity_changes(
+            hass,
+            helper_config_entry_id=helper_config_entry.entry_id,
+            get_helper_entity_id=get_helper_entity_id,
+            set_source_entity_id_or_uuid=set_source_entity_id_or_uuid,
+            source_device_id=source_entity_entry.device_id,
+            source_entity_id_or_uuid=helper_config_entry.options["source"],
+        )
+        return True
+
+    mock_integration(
+        hass,
+        MockModule(
+            HELPER_DOMAIN,
+            async_remove_entry=async_remove_entry,
+            async_setup_entry=async_setup_entry,
+            async_unload_entry=async_unload_entry,
+        ),
+    )
+    mock_platform(hass, f"{HELPER_DOMAIN}.config_flow", None)
+
+
+def track_entity_registry_actions(hass: HomeAssistant, entity_id: str) -> list[str]:
+    """Track entity registry actions for an entity."""
+    events = []
+
+    def add_event(event: Event[er.EventEntityRegistryUpdatedData]) -> None:
+        """Add entity registry updated event to the list."""
+        events.append(event.data["action"])
+
+    async_track_entity_registry_updated_event(hass, entity_id, add_event)
+
+    return events
+
+
+@pytest.mark.parametrize("use_entity_registry_id", [True, False])
+@pytest.mark.usefixtures("mock_helper_flow", "mock_helper_integration")
+async def test_async_handle_source_entity_changes_source_entity_removed(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
+    helper_config_entry: MockConfigEntry,
+    helper_entity_entry: er.RegistryEntry,
+    source_config_entry: ConfigEntry,
+    source_device: dr.DeviceEntry,
+    source_entity_entry: er.RegistryEntry,
+    async_remove_entry: AsyncMock,
+    async_unload_entry: AsyncMock,
+    set_source_entity_id_or_uuid: Mock,
+) -> None:
+    """Test the helper config entry is removed when the source entity is removed."""
+    # Add the helper config entry to the source device
+    device_registry.async_update_device(
+        source_device.id, add_config_entry_id=helper_config_entry.entry_id
+    )
+    # Add another config entry to the source device
+    other_config_entry = MockConfigEntry()
+    other_config_entry.add_to_hass(hass)
+    device_registry.async_update_device(
+        source_device.id, add_config_entry_id=other_config_entry.entry_id
+    )
+
+    assert await hass.config_entries.async_setup(helper_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    # Check preconditions
+    helper_entity_entry = entity_registry.async_get(helper_entity_entry.entity_id)
+    assert helper_entity_entry.device_id == source_entity_entry.device_id
+    source_device = device_registry.async_get(source_device.id)
+    assert helper_config_entry.entry_id in source_device.config_entries
+
+    events = track_entity_registry_actions(hass, helper_entity_entry.entity_id)
+
+    # Remove the source entitys's config entry from the device, this removes the
+    # source entity
+    device_registry.async_update_device(
+        source_device.id, remove_config_entry_id=source_config_entry.entry_id
+    )
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    # Check that the helper config entry is unloaded and removed
+    async_unload_entry.assert_called_once()
+    async_remove_entry.assert_called_once()
+    set_source_entity_id_or_uuid.assert_not_called()
+
+    # Check that the helper config entry is removed from the device
+    source_device = device_registry.async_get(source_device.id)
+    assert helper_config_entry.entry_id not in source_device.config_entries
+
+    # Check that the helper config entry is removed
+    assert helper_config_entry.entry_id not in hass.config_entries.async_entry_ids()
+
+    # Check we got the expected events
+    assert events == ["remove"]
+
+
+@pytest.mark.parametrize("use_entity_registry_id", [True, False])
+@pytest.mark.usefixtures("mock_helper_flow", "mock_helper_integration")
+async def test_async_handle_source_entity_changes_source_entity_removed_from_device(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
+    helper_config_entry: MockConfigEntry,
+    helper_entity_entry: er.RegistryEntry,
+    source_device: dr.DeviceEntry,
+    source_entity_entry: er.RegistryEntry,
+    async_remove_entry: AsyncMock,
+    async_unload_entry: AsyncMock,
+    set_source_entity_id_or_uuid: Mock,
+) -> None:
+    """Test the source entity removed from the source device."""
+    # Add the helper config entry to the source device
+    device_registry.async_update_device(
+        source_device.id, add_config_entry_id=helper_config_entry.entry_id
+    )
+
+    assert await hass.config_entries.async_setup(helper_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    # Check preconditions
+    helper_entity_entry = entity_registry.async_get(helper_entity_entry.entity_id)
+    assert helper_entity_entry.device_id == source_entity_entry.device_id
+
+    source_device = device_registry.async_get(source_device.id)
+    assert helper_config_entry.entry_id in source_device.config_entries
+
+    events = track_entity_registry_actions(hass, helper_entity_entry.entity_id)
+
+    # Remove the source entity from the device
+    entity_registry.async_update_entity(source_entity_entry.entity_id, device_id=None)
+    await hass.async_block_till_done()
+    async_remove_entry.assert_not_called()
+    async_unload_entry.assert_called_once()
+    set_source_entity_id_or_uuid.assert_not_called()
+
+    # Check that the helper config entry is removed from the device
+    source_device = device_registry.async_get(source_device.id)
+    assert helper_config_entry.entry_id not in source_device.config_entries
+
+    # Check that the helper config entry is not removed
+    assert helper_config_entry.entry_id in hass.config_entries.async_entry_ids()
+
+    # Check we got the expected events
+    assert events == ["update"]
+
+
+@pytest.mark.parametrize("use_entity_registry_id", [True, False])
+@pytest.mark.usefixtures("mock_helper_flow", "mock_helper_integration")
+async def test_async_handle_source_entity_changes_source_entity_moved_other_device(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
+    helper_config_entry: MockConfigEntry,
+    helper_entity_entry: er.RegistryEntry,
+    source_config_entry: ConfigEntry,
+    source_device: dr.DeviceEntry,
+    source_entity_entry: er.RegistryEntry,
+    async_remove_entry: AsyncMock,
+    async_unload_entry: AsyncMock,
+    set_source_entity_id_or_uuid: Mock,
+) -> None:
+    """Test the source entity is moved to another device."""
+    # Add the helper config entry to the source device
+    device_registry.async_update_device(
+        source_device.id, add_config_entry_id=helper_config_entry.entry_id
+    )
+
+    # Create another device to move the source entity to
+    source_device_2 = device_registry.async_get_or_create(
+        config_entry_id=source_config_entry.entry_id,
+        connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:FF")},
+    )
+
+    assert await hass.config_entries.async_setup(helper_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    # Check preconditions
+    helper_entity_entry = entity_registry.async_get(helper_entity_entry.entity_id)
+    assert helper_entity_entry.device_id == source_entity_entry.device_id
+
+    source_device = device_registry.async_get(source_device.id)
+    assert helper_config_entry.entry_id in source_device.config_entries
+    source_device_2 = device_registry.async_get(source_device_2.id)
+    assert helper_config_entry.entry_id not in source_device_2.config_entries
+
+    events = track_entity_registry_actions(hass, helper_entity_entry.entity_id)
+
+    # Move the source entity to another device
+    entity_registry.async_update_entity(
+        source_entity_entry.entity_id, device_id=source_device_2.id
+    )
+    await hass.async_block_till_done()
+    async_remove_entry.assert_not_called()
+    async_unload_entry.assert_called_once()
+    set_source_entity_id_or_uuid.assert_not_called()
+
+    # Check that the helper config entry is moved to the other device
+    source_device = device_registry.async_get(source_device.id)
+    assert helper_config_entry.entry_id not in source_device.config_entries
+    source_device_2 = device_registry.async_get(source_device_2.id)
+    assert helper_config_entry.entry_id in source_device_2.config_entries
+
+    # Check that the helper config entry is not removed
+    assert helper_config_entry.entry_id in hass.config_entries.async_entry_ids()
+
+    # Check we got the expected events
+    assert events == ["update"]
+
+
+@pytest.mark.parametrize(
+    ("use_entity_registry_id", "unload_calls", "set_source_entity_id_calls"),
+    [(True, 1, 0), (False, 0, 1)],
+)
+@pytest.mark.usefixtures("mock_helper_flow", "mock_helper_integration")
+async def test_async_handle_source_entity_new_entity_id(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
+    helper_config_entry: MockConfigEntry,
+    helper_entity_entry: er.RegistryEntry,
+    source_device: dr.DeviceEntry,
+    source_entity_entry: er.RegistryEntry,
+    async_remove_entry: AsyncMock,
+    async_unload_entry: AsyncMock,
+    set_source_entity_id_or_uuid: Mock,
+    unload_calls: int,
+    set_source_entity_id_calls: int,
+) -> None:
+    """Test the source entity's entity ID is changed."""
+    # Add the helper config entry to the source device
+    device_registry.async_update_device(
+        source_device.id, add_config_entry_id=helper_config_entry.entry_id
+    )
+
+    assert await hass.config_entries.async_setup(helper_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    # Check preconditions
+    helper_entity_entry = entity_registry.async_get(helper_entity_entry.entity_id)
+    assert helper_entity_entry.device_id == source_entity_entry.device_id
+
+    source_device = device_registry.async_get(source_device.id)
+    assert helper_config_entry.entry_id in source_device.config_entries
+
+    events = track_entity_registry_actions(hass, helper_entity_entry.entity_id)
+
+    # Change the source entity's entity ID
+    entity_registry.async_update_entity(
+        source_entity_entry.entity_id, new_entity_id="sensor.new_entity_id"
+    )
+    await hass.async_block_till_done()
+    async_remove_entry.assert_not_called()
+    assert len(async_unload_entry.mock_calls) == unload_calls
+    assert len(set_source_entity_id_or_uuid.mock_calls) == set_source_entity_id_calls
+
+    # Check that the helper config is still in the device
+    source_device = device_registry.async_get(source_device.id)
+    assert helper_config_entry.entry_id in source_device.config_entries
+
+    # Check that the helper config entry is not removed
+    assert helper_config_entry.entry_id in hass.config_entries.async_entry_ids()
+
+    # Check we got the expected events
+    assert events == []