Do not remove derivative config entry when input sensor is removed (#146506)

* Do not remove derivative config entry when input sensor is removed

* Add comments

* Update homeassistant/helpers/helper_integration.py

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

---------

Co-authored-by: Franck Nijhof <git@frenck.dev>
Co-authored-by: Martin Hjelmare <marhje52@gmail.com>
This commit is contained in:
Erik Montnemery 2025-06-11 11:19:44 +02:00 committed by GitHub
parent 5b4c309170
commit 2afdec4711
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 51 additions and 18 deletions

View File

@ -29,6 +29,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
options={**entry.options, CONF_SOURCE: source_entity_id}, 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)
entity_registry = er.async_get(hass) entity_registry = er.async_get(hass)
entry.async_on_unload( entry.async_on_unload(
async_handle_source_entity_changes( async_handle_source_entity_changes(
@ -42,6 +46,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
hass, entry.options[CONF_SOURCE] hass, entry.options[CONF_SOURCE]
), ),
source_entity_id_or_uuid=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,)) await hass.config_entries.async_forward_entry_setups(entry, (Platform.SENSOR,))

View File

@ -60,6 +60,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
options={**entry.options, CONF_ENTITY_ID: source_entity_id}, options={**entry.options, CONF_ENTITY_ID: source_entity_id},
) )
async def source_entity_removed() -> None:
# The source entity has been removed, we remove the config entry because
# switch_as_x does not allow replacing the wrapped entity.
await hass.config_entries.async_remove(entry.entry_id)
entry.async_on_unload( entry.async_on_unload(
async_handle_source_entity_changes( async_handle_source_entity_changes(
hass, hass,
@ -70,6 +75,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
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],
source_entity_removed=source_entity_removed,
) )
) )
entry.async_on_unload(entry.add_update_listener(config_entry_update_listener)) entry.async_on_unload(entry.add_update_listener(config_entry_update_listener))

View File

@ -62,7 +62,7 @@ def async_device_info_to_link_from_device_id(
def async_remove_stale_devices_links_keep_entity_device( def async_remove_stale_devices_links_keep_entity_device(
hass: HomeAssistant, hass: HomeAssistant,
entry_id: str, entry_id: str,
source_entity_id_or_uuid: str, source_entity_id_or_uuid: str | None,
) -> None: ) -> None:
"""Remove entry_id from all devices except that of source_entity_id_or_uuid. """Remove entry_id from all devices except that of source_entity_id_or_uuid.
@ -73,7 +73,9 @@ def async_remove_stale_devices_links_keep_entity_device(
async_remove_stale_devices_links_keep_current_device( async_remove_stale_devices_links_keep_current_device(
hass=hass, hass=hass,
entry_id=entry_id, entry_id=entry_id,
current_device_id=async_entity_id_to_device_id(hass, source_entity_id_or_uuid), current_device_id=async_entity_id_to_device_id(hass, source_entity_id_or_uuid)
if source_entity_id_or_uuid
else None,
) )

View File

@ -2,7 +2,8 @@
from __future__ import annotations from __future__ import annotations
from collections.abc import Callable from collections.abc import Callable, Coroutine
from typing import Any
from homeassistant.core import CALLBACK_TYPE, Event, HomeAssistant, valid_entity_id from homeassistant.core import CALLBACK_TYPE, Event, HomeAssistant, valid_entity_id
@ -18,6 +19,7 @@ def async_handle_source_entity_changes(
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,
source_entity_removed: Callable[[], Coroutine[Any, Any, None]],
) -> CALLBACK_TYPE: ) -> CALLBACK_TYPE:
"""Handle changes to a helper entity's source entity. """Handle changes to a helper entity's source entity.
@ -34,6 +36,14 @@ def async_handle_source_entity_changes(
- Source entity removed from the device: The helper entity is updated to link - Source entity removed from the device: The helper entity is updated to link
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
ID or UUID, e.g., in the helper config entry options.
:param source_entity_removed: A function which is called when the source entity
is removed. This can be used to clean up any resources related to the source
entity or ask the user to select a new source entity.
""" """
async def async_registry_updated( async def async_registry_updated(
@ -44,7 +54,7 @@ def async_handle_source_entity_changes(
data = event.data data = event.data
if data["action"] == "remove": if data["action"] == "remove":
await hass.config_entries.async_remove(helper_config_entry_id) await source_entity_removed()
if data["action"] != "update": if data["action"] != "update":
return return

View File

@ -268,17 +268,17 @@ async def test_async_handle_source_entity_changes_source_entity_removed(
) )
await hass.async_block_till_done() await hass.async_block_till_done()
await hass.async_block_till_done() await hass.async_block_till_done()
mock_unload_entry.assert_called_once() mock_unload_entry.assert_not_called()
# Check that the derivative config entry is removed from the device # Check that the derivative config entry is removed from the device
sensor_device = device_registry.async_get(sensor_device.id) sensor_device = device_registry.async_get(sensor_device.id)
assert derivative_config_entry.entry_id not in sensor_device.config_entries assert derivative_config_entry.entry_id not in sensor_device.config_entries
# Check that the derivative config entry is removed # Check that the derivative config entry is not removed
assert derivative_config_entry.entry_id not in hass.config_entries.async_entry_ids() assert derivative_config_entry.entry_id in hass.config_entries.async_entry_ids()
# Check we got the expected events # Check we got the expected events
assert events == ["remove"] assert events == ["update"]
async def test_async_handle_source_entity_changes_source_entity_removed_from_device( async def test_async_handle_source_entity_changes_source_entity_removed_from_device(

View File

@ -132,11 +132,17 @@ def async_unload_entry() -> AsyncMock:
@pytest.fixture @pytest.fixture
def set_source_entity_id_or_uuid() -> AsyncMock: def set_source_entity_id_or_uuid() -> Mock:
"""Fixture to mock async_unload_entry.""" """Fixture to mock set_source_entity_id_or_uuid."""
return Mock() return Mock()
@pytest.fixture
def source_entity_removed() -> AsyncMock:
"""Fixture to mock source_entity_removed."""
return AsyncMock()
@pytest.fixture @pytest.fixture
def mock_helper_integration( def mock_helper_integration(
hass: HomeAssistant, hass: HomeAssistant,
@ -146,6 +152,7 @@ def mock_helper_integration(
async_remove_entry: AsyncMock, async_remove_entry: AsyncMock,
async_unload_entry: AsyncMock, async_unload_entry: AsyncMock,
set_source_entity_id_or_uuid: Mock, set_source_entity_id_or_uuid: Mock,
source_entity_removed: AsyncMock,
) -> None: ) -> None:
"""Mock the helper integration.""" """Mock the helper integration."""
@ -164,6 +171,7 @@ def mock_helper_integration(
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"],
source_entity_removed=source_entity_removed,
) )
return True return True
@ -206,6 +214,7 @@ async def test_async_handle_source_entity_changes_source_entity_removed(
async_remove_entry: AsyncMock, async_remove_entry: AsyncMock,
async_unload_entry: AsyncMock, async_unload_entry: AsyncMock,
set_source_entity_id_or_uuid: Mock, set_source_entity_id_or_uuid: Mock,
source_entity_removed: AsyncMock,
) -> None: ) -> None:
"""Test the helper config entry is removed when the source entity is removed.""" """Test the helper config entry is removed when the source entity is removed."""
# Add the helper config entry to the source device # Add the helper config entry to the source device
@ -238,20 +247,21 @@ async def test_async_handle_source_entity_changes_source_entity_removed(
await hass.async_block_till_done() await hass.async_block_till_done()
await hass.async_block_till_done() await hass.async_block_till_done()
# Check that the helper config entry is unloaded and removed # Check that the source_entity_removed callback was called
async_unload_entry.assert_called_once() source_entity_removed.assert_called_once()
async_remove_entry.assert_called_once() async_unload_entry.assert_not_called()
async_remove_entry.assert_not_called()
set_source_entity_id_or_uuid.assert_not_called() set_source_entity_id_or_uuid.assert_not_called()
# Check that the helper config entry is removed from the device # Check that the helper config entry is not removed from the device
source_device = device_registry.async_get(source_device.id) source_device = device_registry.async_get(source_device.id)
assert helper_config_entry.entry_id not in source_device.config_entries assert helper_config_entry.entry_id in source_device.config_entries
# Check that the helper config entry is removed # Check that the helper config entry is not removed
assert helper_config_entry.entry_id not in hass.config_entries.async_entry_ids() assert helper_config_entry.entry_id in hass.config_entries.async_entry_ids()
# Check we got the expected events # Check we got the expected events
assert events == ["remove"] assert events == []
@pytest.mark.parametrize("use_entity_registry_id", [True, False]) @pytest.mark.parametrize("use_entity_registry_id", [True, False])