Simplify helper_integration.async_handle_source_entity_changes (#146516)

This commit is contained in:
Erik Montnemery 2025-06-11 12:27:51 +02:00 committed by GitHub
parent f908e0cf4d
commit c4c8f88765
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 7 additions and 27 deletions

View File

@ -2,19 +2,15 @@
from __future__ import annotations from __future__ import annotations
from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
from homeassistant.config_entries import ConfigEntry from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_SOURCE, Platform from homeassistant.const import CONF_SOURCE, Platform
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.device import ( from homeassistant.helpers.device import (
async_entity_id_to_device_id, async_entity_id_to_device_id,
async_remove_stale_devices_links_keep_entity_device, 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
from .const import DOMAIN
async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Set up Derivative from a config entry.""" """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. # 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) async_remove_stale_devices_links_keep_entity_device(hass, entry.entry_id, None)
entity_registry = er.async_get(hass)
entry.async_on_unload( entry.async_on_unload(
async_handle_source_entity_changes( async_handle_source_entity_changes(
hass, hass,
helper_config_entry_id=entry.entry_id, 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, set_source_entity_id_or_uuid=set_source_entity_id_or_uuid,
source_device_id=async_entity_id_to_device_id( source_device_id=async_entity_id_to_device_id(
hass, entry.options[CONF_SOURCE] hass, entry.options[CONF_SOURCE]

View File

@ -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 import device_registry as dr, entity_registry as er
from homeassistant.helpers.helper_integration import async_handle_source_entity_changes 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__) _LOGGER = logging.getLogger(__name__)
@ -69,9 +69,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
async_handle_source_entity_changes( async_handle_source_entity_changes(
hass, hass,
helper_config_entry_id=entry.entry_id, 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, set_source_entity_id_or_uuid=set_source_entity_id_or_uuid,
source_device_id=async_add_to_device(hass, entry, entity_id), source_device_id=async_add_to_device(hass, entry, entity_id),
source_entity_id_or_uuid=entry.options[CONF_ENTITY_ID], source_entity_id_or_uuid=entry.options[CONF_ENTITY_ID],

View File

@ -15,7 +15,6 @@ def async_handle_source_entity_changes(
hass: HomeAssistant, hass: HomeAssistant,
*, *,
helper_config_entry_id: str, helper_config_entry_id: str,
get_helper_entity_id: Callable[[], str | None],
set_source_entity_id_or_uuid: Callable[[str], None], set_source_entity_id_or_uuid: Callable[[str], None],
source_device_id: str | None, source_device_id: str | None,
source_entity_id_or_uuid: str, 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 to no device, and the helper config entry removed from the old device. Then
the helper config entry is reloaded. 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 :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. ID or UUID, e.g., in the helper config entry options.
:param source_entity_removed: A function which is called when the source entity :param source_entity_removed: A function which is called when the source entity
@ -81,13 +78,14 @@ def async_handle_source_entity_changes(
return return
# The source entity has been moved to a different device, update the helper # 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 # entities to link to the new device and the helper device to include the
# the helper config entry # helper config entry
helper_entity_id = get_helper_entity_id() for helper_entity in entity_registry.entities.get_entries_for_config_entry_id(
if helper_entity_id: helper_config_entry_id
):
# Update the helper entity to link to the new device (or no device) # Update the helper entity to link to the new device (or no device)
entity_registry.async_update_entity( 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: if source_entity_entry.device_id is not None:

View File

@ -156,18 +156,11 @@ def mock_helper_integration(
) -> None: ) -> None:
"""Mock the helper integration.""" """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: async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Mock setup entry.""" """Mock setup entry."""
async_handle_source_entity_changes( async_handle_source_entity_changes(
hass, hass,
helper_config_entry_id=helper_config_entry.entry_id, 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, set_source_entity_id_or_uuid=set_source_entity_id_or_uuid,
source_device_id=source_entity_entry.device_id, source_device_id=source_entity_entry.device_id,
source_entity_id_or_uuid=helper_config_entry.options["source"], source_entity_id_or_uuid=helper_config_entry.options["source"],