diff --git a/homeassistant/components/derivative/__init__.py b/homeassistant/components/derivative/__init__.py index 6d539817875..0806a8f824d 100644 --- a/homeassistant/components/derivative/__init__.py +++ b/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 async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: # 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] diff --git a/homeassistant/components/switch_as_x/__init__.py b/homeassistant/components/switch_as_x/__init__.py index 6bd2c4e6482..c77eda9b294 100644 --- a/homeassistant/components/switch_as_x/__init__.py +++ b/homeassistant/components/switch_as_x/__init__.py @@ -13,7 +13,7 @@ from homeassistant.core import HomeAssistant, callback 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 .const import CONF_INVERT, CONF_TARGET_DOMAIN _LOGGER = logging.getLogger(__name__) @@ -69,9 +69,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: 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], diff --git a/homeassistant/helpers/helper_integration.py b/homeassistant/helpers/helper_integration.py index 37aa246178e..61bb0bcd45d 100644 --- a/homeassistant/helpers/helper_integration.py +++ b/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 @@ def async_handle_source_entity_changes( 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: diff --git a/tests/helpers/test_helper_integration.py b/tests/helpers/test_helper_integration.py index 12433894dc7..47f1b62feb7 100644 --- a/tests/helpers/test_helper_integration.py +++ b/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"],