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

Co-authored-by: Artur Pragacz <[email protected]>

Author: emontnemery

homeassistant/components/derivative/__init__.py

@@ -11,14 +11,18 @@
     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__)
 
 
 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 @@ def set_source_entity_id_or_uuid(source_entity_id: str) -> None:
             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 @@
     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_registry import DeviceInfo
+from homeassistant.helpers.device import async_entity_id_to_device
 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 @@ def __init__(
         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
-      is removed, and the helper entity is cleaned up.
+    - Entity removal: If the source entity is removed:
+      - 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 @@ async def async_registry_updated(
 
         data = event.data
         if data["action"] == "remove":
-            await source_entity_removed()
+            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