Simplify helper_integration.async_handle_source_entity_changes (home-assistant#146516)

Author: emontnemery

homeassistant/components/derivative/__init__.py

@@ -2,19 +2,15 @@
 
 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:
     """Set up Derivative from a config entry."""
@@ -33,14 +29,10 @@ 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)
 
-    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]

homeassistant/components/switch_as_x/__init__.py

@@ -13,10 +13,7 @@
 from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.helpers.helper_integration import async_handle_source_entity_changes
 
-from .const import CONF_INVERT, CONF_TARGET_DOMAIN, DOMAIN
-from .light import LightSwitch
-
-__all__ = ["LightSwitch"]
+from .const import CONF_INVERT, CONF_TARGET_DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -72,9 +69,6 @@ async def source_entity_removed() -> None:
         async_handle_source_entity_changes(
             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],

homeassistant/helpers/helper_integration.py

@@ -15,7 +15,6 @@ 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,
@@ -37,8 +36,6 @@ def async_handle_source_entity_changes(
       to no device, and the helper config entry removed from the old device. Then
       the helper config entry is reloaded.
 
-    :param get_helper_entity: A function which returns the helper entity's entity ID,
-        or None if the helper entity does not exist.
     :param set_source_entity_id_or_uuid: A function which updates the source entity
         ID or UUID, e.g., in the helper config entry options.
     :param source_entity_removed: A function which is called when the source entity
@@ -81,13 +78,14 @@ async def async_registry_updated(
             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:
+        # entities to link to the new device and the helper device to include the
+        # helper config entry
+        for helper_entity 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_id, device_id=source_entity_entry.device_id
+                helper_entity.entity_id, device_id=source_entity_entry.device_id
             )
 
         if source_entity_entry.device_id is not None:

tests/helpers/test_helper_integration.py

@@ -156,18 +156,11 @@ def mock_helper_integration(
 ) -> 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"],