Do not add derivative config entry to source device (#148674)

Co-authored-by: Artur Pragacz <49985303+arturpragacz@users.noreply.github.com>

homeassistant/components/derivative/__init__.py

@@ -11,7 +11,10 @@ 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 homeassistant.helpers.helper_integration import (
+    async_handle_source_entity_changes,
+    async_remove_helper_config_entry_from_source_device,
+)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -19,6 +22,7 @@ _LOGGER = logging.getLogger(__name__)
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Derivative from a config entry."""
 
+    # This can be removed in HA Core 2026.2
     async_remove_stale_devices_links_keep_entity_device(
         hass, entry.entry_id, entry.options[CONF_SOURCE]
     )
@@ -29,20 +33,16 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             options={**entry.options, CONF_SOURCE: source_entity_id},
         )
 
-    async def source_entity_removed() -> None:
-        # The source entity has been removed, we need to clean the device links.
-        async_remove_stale_devices_links_keep_entity_device(hass, entry.entry_id, None)
-
     entry.async_on_unload(
         async_handle_source_entity_changes(
             hass,
+            add_helper_config_entry_to_device=False,
             helper_config_entry_id=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],
-            source_entity_removed=source_entity_removed,
         )
     )
     await hass.config_entries.async_forward_entry_setups(entry, (Platform.SENSOR,))
@@ -85,6 +85,20 @@ async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
                 config_entry, options=new_options, version=1, minor_version=2
             )
 
+        if config_entry.minor_version < 3:
+            # Remove the derivative config entry from the source device
+            if source_device_id := async_entity_id_to_device_id(
+                hass, config_entry.options[CONF_SOURCE]
+            ):
+                async_remove_helper_config_entry_from_source_device(
+                    hass,
+                    helper_config_entry_id=config_entry.entry_id,
+                    source_device_id=source_device_id,
+                )
+            hass.config_entries.async_update_entry(
+                config_entry, version=1, minor_version=3
+            )
+
     _LOGGER.debug(
         "Migration to configuration version %s.%s successful",
         config_entry.version,

homeassistant/components/derivative/config_flow.py

@@ -142,7 +142,7 @@ class ConfigFlowHandler(SchemaConfigFlowHandler, domain=DOMAIN):
     options_flow = OPTIONS_FLOW
 
     VERSION = 1
-    MINOR_VERSION = 2
+    MINOR_VERSION = 3
 
     def async_config_entry_title(self, options: Mapping[str, Any]) -> str:
         """Return config entry title."""

homeassistant/components/derivative/sensor.py

@@ -34,8 +34,7 @@ from homeassistant.core import (
     callback,
 )
 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 import async_entity_id_to_device
-from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import (
     AddConfigEntryEntitiesCallback,
     AddEntitiesCallback,
@@ -118,17 +117,13 @@ async def async_setup_entry(
         registry, config_entry.options[CONF_SOURCE]
     )
 
-    device_info = async_device_info_to_link_from_entity(
-        hass,
-        source_entity_id,
-    )
-
     if max_sub_interval_dict := config_entry.options.get(CONF_MAX_SUB_INTERVAL, None):
         max_sub_interval = cv.time_period(max_sub_interval_dict)
     else:
         max_sub_interval = None
 
     derivative_sensor = DerivativeSensor(
+        hass,
         name=config_entry.title,
         round_digits=int(config_entry.options[CONF_ROUND_DIGITS]),
         source_entity=source_entity_id,
@@ -137,7 +132,6 @@ async def async_setup_entry(
         unit_of_measurement=None,
         unit_prefix=config_entry.options.get(CONF_UNIT_PREFIX),
         unit_time=config_entry.options[CONF_UNIT_TIME],
-        device_info=device_info,
         max_sub_interval=max_sub_interval,
     )
 
@@ -152,6 +146,7 @@ async def async_setup_platform(
 ) -> None:
     """Set up the derivative sensor."""
     derivative = DerivativeSensor(
+        hass,
         name=config.get(CONF_NAME),
         round_digits=config[CONF_ROUND_DIGITS],
         source_entity=config[CONF_SOURCE],
@@ -174,6 +169,7 @@ class DerivativeSensor(RestoreSensor, SensorEntity):
 
     def __init__(
         self,
+        hass: HomeAssistant,
         *,
         name: str | None,
         round_digits: int,
@@ -184,11 +180,13 @@ class DerivativeSensor(RestoreSensor, SensorEntity):
         unit_time: UnitOfTime,
         max_sub_interval: timedelta | None,
         unique_id: str | None,
-        device_info: DeviceInfo | None = None,
     ) -> None:
         """Initialize the derivative sensor."""
         self._attr_unique_id = unique_id
-        self._attr_device_info = device_info
+        self.device_entry = async_entity_id_to_device(
+            hass,
+            source_entity,
+        )
         self._sensor_source_id = source_entity
         self._round_digits = round_digits
         self._attr_native_value = round(Decimal(0), round_digits)

homeassistant/helpers/device.py

@@ -21,6 +21,19 @@ def async_entity_id_to_device_id(
     return entity.device_id
 
 
+@callback
+def async_entity_id_to_device(
+    hass: HomeAssistant,
+    entity_id_or_uuid: str,
+) -> dr.DeviceEntry | None:
+    """Resolve the device entry for the entity id or entity uuid."""
+
+    if (device_id := async_entity_id_to_device_id(hass, entity_id_or_uuid)) is None:
+        return None
+
+    return dr.async_get(hass).async_get(device_id)
+
+
 @callback
 def async_device_info_to_link_from_entity(
     hass: HomeAssistant,

homeassistant/helpers/helper_integration.py

@@ -19,13 +19,15 @@ def async_handle_source_entity_changes(
     set_source_entity_id_or_uuid: Callable[[str], None],
     source_device_id: str | None,
     source_entity_id_or_uuid: str,
-    source_entity_removed: Callable[[], Coroutine[Any, Any, None]],
+    source_entity_removed: Callable[[], Coroutine[Any, Any, None]] | None = None,
 ) -> 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
+    - Entity removal: If the source entity is removed:
-      is removed, and the helper entity is cleaned up.
+      - If source_entity_removed is provided, it is called to handle the removal.
+      - If source_entity_removed is not provided, The helper entity is updated to
+      not link to any device.
     - 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
@@ -52,7 +54,18 @@ def async_handle_source_entity_changes(
 
         data = event.data
         if data["action"] == "remove":
+            if source_entity_removed:
                 await source_entity_removed()
+            else:
+                for (
+                    helper_entity_entry
+                ) in entity_registry.entities.get_entries_for_config_entry_id(
+                    helper_config_entry_id
+                ):
+                    # Update the helper entity to link to the new device (or no device)
+                    entity_registry.async_update_entity(
+                        helper_entity_entry.entity_id, device_id=None
+                    )
 
         if data["action"] != "update":
             return

tests/components/derivative/test_init.py

@@ -8,7 +8,7 @@ from homeassistant.components import derivative
 from homeassistant.components.derivative.config_flow import ConfigFlowHandler
 from homeassistant.components.derivative.const import DOMAIN
 from homeassistant.config_entries import ConfigEntry, ConfigEntryState
-from homeassistant.core import Event, HomeAssistant
+from homeassistant.core import Event, 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
 
@@ -82,6 +82,7 @@ def track_entity_registry_actions(hass: HomeAssistant, entity_id: str) -> list[s
     """Track entity registry actions for an entity."""
     events = []
 
+    @callback
     def add_event(event: Event[er.EventEntityRegistryUpdatedData]) -> None:
         """Add entity registry updated event to the list."""
         events.append(event.data["action"])
@@ -214,7 +215,7 @@ async def test_device_cleaning(
     devices_before_reload = device_registry.devices.get_devices_for_config_entry_id(
         derivative_config_entry.entry_id
     )
-    assert len(devices_before_reload) == 3
+    assert len(devices_before_reload) == 2
 
     # Config entry reload
     await hass.config_entries.async_reload(derivative_config_entry.entry_id)
@@ -229,7 +230,7 @@ async def test_device_cleaning(
     devices_after_reload = device_registry.devices.get_devices_for_config_entry_id(
         derivative_config_entry.entry_id
     )
-    assert len(devices_after_reload) == 1
+    assert len(devices_after_reload) == 0
 
 
 async def test_async_handle_source_entity_changes_source_entity_removed(
@@ -240,6 +241,54 @@ async def test_async_handle_source_entity_changes_source_entity_removed(
     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."""
+    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 not 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_not_called()
+
+    # Check that the entity is no longer linked to the source device
+    derivative_entity_entry = entity_registry.async_get("sensor.my_derivative")
+    assert derivative_entity_entry.device_id is None
+
+    # Check that the device is removed
+    assert not device_registry.async_get(sensor_device.id)
+
+    # 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_removed_shared_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 derivative config entry is removed when the source entity is removed."""
     # Add another config entry to the sensor device
@@ -256,7 +305,7 @@ async def test_async_handle_source_entity_changes_source_entity_removed(
     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
+    assert derivative_config_entry.entry_id not in sensor_device.config_entries
 
     events = track_entity_registry_actions(hass, derivative_entity_entry.entity_id)
 
@@ -273,7 +322,11 @@ async def test_async_handle_source_entity_changes_source_entity_removed(
         await hass.async_block_till_done()
     mock_unload_entry.assert_not_called()
 
-    # Check that the derivative config entry is removed from the device
+    # Check that the entity is no longer linked to the source device
+    derivative_entity_entry = entity_registry.async_get("sensor.my_derivative")
+    assert derivative_entity_entry.device_id is None
+
+    # Check that the derivative config entry is not in the device
     sensor_device = device_registry.async_get(sensor_device.id)
     assert derivative_config_entry.entry_id not in sensor_device.config_entries
 
@@ -300,7 +353,7 @@ async def test_async_handle_source_entity_changes_source_entity_removed_from_dev
     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
+    assert derivative_config_entry.entry_id not in sensor_device.config_entries
 
     events = track_entity_registry_actions(hass, derivative_entity_entry.entity_id)
 
@@ -315,7 +368,11 @@ async def test_async_handle_source_entity_changes_source_entity_removed_from_dev
         await hass.async_block_till_done()
     mock_unload_entry.assert_called_once()
 
-    # Check that the derivative config entry is removed from the device
+    # Check that the entity is no longer linked to the source device
+    derivative_entity_entry = entity_registry.async_get("sensor.my_derivative")
+    assert derivative_entity_entry.device_id is None
+
+    # Check that the derivative config entry is not in the device
     sensor_device = device_registry.async_get(sensor_device.id)
     assert derivative_config_entry.entry_id not in sensor_device.config_entries
 
@@ -348,7 +405,7 @@ async def test_async_handle_source_entity_changes_source_entity_moved_other_devi
     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
+    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 not in sensor_device_2.config_entries
 
@@ -365,11 +422,15 @@ async def test_async_handle_source_entity_changes_source_entity_moved_other_devi
         await hass.async_block_till_done()
     mock_unload_entry.assert_called_once()
 
-    # Check that the derivative config entry is moved to the other device
+    # Check that the entity is linked to the other device
+    derivative_entity_entry = entity_registry.async_get("sensor.my_derivative")
+    assert derivative_entity_entry.device_id == sensor_device_2.id
+
+    # Check that the derivative config entry is not in any of the devices
     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
+    assert derivative_config_entry.entry_id not 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()
@@ -394,7 +455,7 @@ async def test_async_handle_source_entity_new_entity_id(
     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
+    assert derivative_config_entry.entry_id not in sensor_device.config_entries
 
     events = track_entity_registry_actions(hass, derivative_entity_entry.entity_id)
 
@@ -412,9 +473,9 @@ async def test_async_handle_source_entity_new_entity_id(
     # 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
+    # Check that the helper config is not in the device
     sensor_device = device_registry.async_get(sensor_device.id)
-    assert derivative_config_entry.entry_id in sensor_device.config_entries
+    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()
@@ -431,7 +492,7 @@ async def test_async_handle_source_entity_new_entity_id(
         ({"unit_prefix": "none"}, None),
     ],
 )
-async def test_migration(hass: HomeAssistant, unit_prefix, expect_prefix) -> None:
+async def test_migration_1_1(hass: HomeAssistant, unit_prefix, expect_prefix) -> None:
     """Test migration from v1.1 deletes "none" unit_prefix."""
 
     config_entry = MockConfigEntry(
@@ -457,6 +518,60 @@ async def test_migration(hass: HomeAssistant, unit_prefix, expect_prefix) -> Non
     assert config_entry.options["unit_time"] == "min"
     assert config_entry.options.get("unit_prefix") == expect_prefix
 
+    assert config_entry.version == 1
+    assert config_entry.minor_version == 3
+
+
+async def test_migration_1_2(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
+    sensor_device: dr.DeviceEntry,
+    sensor_entity_entry: er.RegistryEntry,
+) -> None:
+    """Test migration from v1.2 removes derivative config entry from device."""
+
+    derivative_config_entry = MockConfigEntry(
+        data={},
+        domain=DOMAIN,
+        options={
+            "name": "My derivative",
+            "round": 1.0,
+            "source": "sensor.test_unique",
+            "time_window": {"seconds": 0.0},
+            "unit_prefix": "k",
+            "unit_time": "min",
+        },
+        title="My derivative",
+        version=1,
+        minor_version=2,
+    )
+    derivative_config_entry.add_to_hass(hass)
+
+    # Add the helper config entry to the device
+    device_registry.async_update_device(
+        sensor_device.id, add_config_entry_id=derivative_config_entry.entry_id
+    )
+
+    # Check preconditions
+    sensor_device = device_registry.async_get(sensor_device.id)
+    assert derivative_config_entry.entry_id in sensor_device.config_entries
+
+    await hass.config_entries.async_setup(derivative_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert derivative_config_entry.state is ConfigEntryState.LOADED
+
+    # Check that the helper config entry is removed from the device and the helper
+    # entity is linked to the source device
+    sensor_device = device_registry.async_get(sensor_device.id)
+    assert derivative_config_entry.entry_id not in sensor_device.config_entries
+    derivative_entity_entry = entity_registry.async_get("sensor.my_derivative")
+    assert derivative_entity_entry.device_id == sensor_entity_entry.device_id
+
+    assert derivative_config_entry.version == 1
+    assert derivative_config_entry.minor_version == 3
+
 
 async def test_migration_from_future_version(
     hass: HomeAssistant,

tests/helpers/test_device.py

@@ -8,6 +8,7 @@ from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.helpers.device import (
     async_device_info_to_link_from_device_id,
     async_device_info_to_link_from_entity,
+    async_entity_id_to_device,
     async_entity_id_to_device_id,
     async_remove_stale_devices_links_keep_current_device,
     async_remove_stale_devices_links_keep_entity_device,
@@ -16,12 +17,12 @@ from homeassistant.helpers.device import (
 from tests.common import MockConfigEntry
 
 
-async def test_entity_id_to_device_id(
+async def test_entity_id_to_device_device_id(
     hass: HomeAssistant,
     device_registry: dr.DeviceRegistry,
     entity_registry: er.EntityRegistry,
 ) -> None:
-    """Test returning an entity's device ID."""
+    """Test returning an entity's device / device ID."""
     config_entry = MockConfigEntry(domain="my")
     config_entry.add_to_hass(hass)
 
@@ -48,6 +49,41 @@ async def test_entity_id_to_device_id(
         entity_id_or_uuid=entity.entity_id,
     )
     assert device_id == device.id
+    assert (
+        async_entity_id_to_device(
+            hass,
+            entity_id_or_uuid=entity.entity_id,
+        )
+        == device
+    )
+
+    assert (
+        async_entity_id_to_device_id(
+            hass,
+            entity_id_or_uuid="unknown.entity_id",
+        )
+        is None
+    )
+    assert (
+        async_entity_id_to_device(
+            hass,
+            entity_id_or_uuid="unknown.entity_id",
+        )
+        is None
+    )
+
+    device_id = async_entity_id_to_device_id(
+        hass,
+        entity_id_or_uuid=entity.id,
+    )
+    assert device_id == device.id
+    assert (
+        async_entity_id_to_device(
+            hass,
+            entity_id_or_uuid=entity.id,
+        )
+        == device
+    )
 
     with pytest.raises(vol.Invalid):
         async_entity_id_to_device_id(
@@ -55,6 +91,12 @@ async def test_entity_id_to_device_id(
             entity_id_or_uuid="unknown_uuid",
         )
 
+    with pytest.raises(vol.Invalid):
+        async_entity_id_to_device(
+            hass,
+            entity_id_or_uuid="unknown_uuid",
+        )
+
 
 async def test_device_info_to_link(
     hass: HomeAssistant,

tests/helpers/test_helper_integration.py

@@ -155,7 +155,7 @@ def mock_helper_integration(
     async_remove_entry: AsyncMock,
     async_unload_entry: AsyncMock,
     set_source_entity_id_or_uuid: Mock,
-    source_entity_removed: AsyncMock,
+    source_entity_removed: AsyncMock | None,
 ) -> None:
     """Mock the helper integration."""
 
@@ -197,7 +197,9 @@ def track_entity_registry_actions(hass: HomeAssistant, entity_id: str) -> list[s
     return events
 
 
-def listen_entity_registry_events(hass: HomeAssistant) -> list[str]:
+def listen_entity_registry_events(
+    hass: HomeAssistant,
+) -> list[er.EventEntityRegistryUpdatedData]:
     """Track entity registry actions for an entity."""
     events: list[er.EventEntityRegistryUpdatedData] = []
 
@@ -211,6 +213,7 @@ def listen_entity_registry_events(hass: HomeAssistant) -> list[str]:
     return events
 
 
+@pytest.mark.parametrize("source_entity_removed", [None])
 @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(
@@ -225,6 +228,70 @@ async def test_async_handle_source_entity_changes_source_entity_removed(
     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 entity is not linked to the source device anymore
+    helper_entity_entry = entity_registry.async_get(helper_entity_entry.entity_id)
+    assert helper_entity_entry.device_id is None
+    async_unload_entry.assert_not_called()
+    async_remove_entry.assert_not_called()
+    set_source_entity_id_or_uuid.assert_not_called()
+
+    # Check that the helper config entry is not removed from 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 == ["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_removed_custom_handler(
+    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,
     source_entity_removed: AsyncMock,
 ) -> None:
     """Test the helper config entry is removed when the source entity is removed."""