diff --git a/homeassistant/components/derivative/__init__.py b/homeassistant/components/derivative/__init__.py index 5117663f3c5..5eb499b0efd 100644 --- a/homeassistant/components/derivative/__init__.py +++ b/homeassistant/components/derivative/__init__.py @@ -2,12 +2,18 @@ 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: @@ -17,6 +23,27 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: hass, entry.entry_id, entry.options[CONF_SOURCE] ) + def set_source_entity_id_or_uuid(source_entity_id: str) -> None: + hass.config_entries.async_update_entry( + entry, + options={**entry.options, CONF_SOURCE: source_entity_id}, + ) + + 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] + ), + source_entity_id_or_uuid=entry.options[CONF_SOURCE], + ) + ) await hass.config_entries.async_forward_entry_setups(entry, (Platform.SENSOR,)) entry.async_on_unload(entry.add_update_listener(config_entry_update_listener)) return True diff --git a/homeassistant/components/switch_as_x/__init__.py b/homeassistant/components/switch_as_x/__init__.py index 2bc0710d08f..6f21b032da3 100644 --- a/homeassistant/components/switch_as_x/__init__.py +++ b/homeassistant/components/switch_as_x/__init__.py @@ -9,9 +9,9 @@ import voluptuous as vol from homeassistant.components.homeassistant import exposed_entities from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_ENTITY_ID -from homeassistant.core import Event, HomeAssistant, callback, valid_entity_id +from homeassistant.core import HomeAssistant, callback from homeassistant.helpers import device_registry as dr, entity_registry as er -from homeassistant.helpers.event import async_track_entity_registry_updated_event +from homeassistant.helpers.helper_integration import async_handle_source_entity_changes from .const import CONF_INVERT, CONF_TARGET_DOMAIN, DOMAIN @@ -42,7 +42,6 @@ def async_add_to_device( async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up a config entry.""" entity_registry = er.async_get(hass) - device_registry = dr.async_get(hass) try: entity_id = er.async_validate_entity_id( entity_registry, entry.options[CONF_ENTITY_ID] @@ -55,72 +54,26 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: ) return False - async def async_registry_updated( - event: Event[er.EventEntityRegistryUpdatedData], - ) -> None: - """Handle entity registry update.""" - data = event.data - if data["action"] == "remove": - await hass.config_entries.async_remove(entry.entry_id) - - if data["action"] != "update": - return - - if "entity_id" in data["changes"]: - # Entity_id changed, update or reload the config entry - if valid_entity_id(entry.options[CONF_ENTITY_ID]): - # If the entity is pointed to by an entity ID, update the entry - hass.config_entries.async_update_entry( - entry, - options={**entry.options, CONF_ENTITY_ID: data["entity_id"]}, - ) - else: - await hass.config_entries.async_reload(entry.entry_id) - - if device_id and "device_id" in data["changes"]: - # Handle the wrapped switch being moved to a different device or removed - # from the device - if ( - not (entity_entry := entity_registry.async_get(data[CONF_ENTITY_ID])) - or not device_registry.async_get(device_id) - or entity_entry.device_id == device_id - ): - # No need to do any cleanup - return - - # The wrapped switch has been moved to a different device, update the - # switch_as_x entity and the device entry to include our config entry - switch_as_x_entity_id = entity_registry.async_get_entity_id( - entry.options[CONF_TARGET_DOMAIN], DOMAIN, entry.entry_id - ) - if switch_as_x_entity_id: - # Update the switch_as_x entity to point to the new device (or no device) - entity_registry.async_update_entity( - switch_as_x_entity_id, device_id=entity_entry.device_id - ) - - if entity_entry.device_id is not None: - device_registry.async_update_device( - entity_entry.device_id, add_config_entry_id=entry.entry_id - ) - - device_registry.async_update_device( - device_id, remove_config_entry_id=entry.entry_id - ) - - # Reload the config entry so the switch_as_x entity is recreated with - # correct device info - await hass.config_entries.async_reload(entry.entry_id) + def set_source_entity_id_or_uuid(source_entity_id: str) -> None: + hass.config_entries.async_update_entry( + entry, + options={**entry.options, CONF_ENTITY_ID: source_entity_id}, + ) entry.async_on_unload( - async_track_entity_registry_updated_event( - hass, entity_id, async_registry_updated + 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], ) ) entry.async_on_unload(entry.add_update_listener(config_entry_update_listener)) - device_id = async_add_to_device(hass, entry, entity_id) - await hass.config_entries.async_forward_entry_setups( entry, (entry.options[CONF_TARGET_DOMAIN],) ) diff --git a/homeassistant/helpers/helper_integration.py b/homeassistant/helpers/helper_integration.py new file mode 100644 index 00000000000..4f39ef4c843 --- /dev/null +++ b/homeassistant/helpers/helper_integration.py @@ -0,0 +1,105 @@ +"""Helpers for helper integrations.""" + +from __future__ import annotations + +from collections.abc import Callable + +from homeassistant.core import CALLBACK_TYPE, Event, HomeAssistant, valid_entity_id + +from . import device_registry as dr, entity_registry as er +from .event import async_track_entity_registry_updated_event + + +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, +) -> 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 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 + is reloaded. + - Source entity moved to another device: The helper entity is updated to link + to the new device, and the helper config entry removed from the old device + and added to the new device. Then the helper config entry is reloaded. + - 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 + the helper config entry is reloaded. + """ + + async def async_registry_updated( + event: Event[er.EventEntityRegistryUpdatedData], + ) -> None: + """Handle entity registry update.""" + nonlocal source_device_id + + data = event.data + if data["action"] == "remove": + await hass.config_entries.async_remove(helper_config_entry_id) + + if data["action"] != "update": + return + + if "entity_id" in data["changes"]: + # Entity_id changed, update or reload the config entry + if valid_entity_id(source_entity_id_or_uuid): + # If the entity is pointed to by an entity ID, update the entry + set_source_entity_id_or_uuid(data["entity_id"]) + else: + await hass.config_entries.async_reload(helper_config_entry_id) + + if not source_device_id or "device_id" not in data["changes"]: + return + + # Handle the source entity being moved to a different device or removed + # from the device + if ( + not (source_entity_entry := entity_registry.async_get(data["entity_id"])) + or not device_registry.async_get(source_device_id) + or source_entity_entry.device_id == source_device_id + ): + # No need to do any cleanup + 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: + # 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 + ) + + if source_entity_entry.device_id is not None: + device_registry.async_update_device( + source_entity_entry.device_id, + add_config_entry_id=helper_config_entry_id, + ) + + device_registry.async_update_device( + source_device_id, remove_config_entry_id=helper_config_entry_id + ) + source_device_id = source_entity_entry.device_id + + # Reload the config entry so the helper entity is recreated with + # correct device info + await hass.config_entries.async_reload(helper_config_entry_id) + + device_registry = dr.async_get(hass) + entity_registry = er.async_get(hass) + source_entity_id = er.async_validate_entity_id( + entity_registry, source_entity_id_or_uuid + ) + return async_track_entity_registry_updated_event( + hass, source_entity_id, async_registry_updated + ) diff --git a/tests/components/derivative/test_init.py b/tests/components/derivative/test_init.py index 32802080e39..f75d5940da7 100644 --- a/tests/components/derivative/test_init.py +++ b/tests/components/derivative/test_init.py @@ -1,23 +1,103 @@ """Test the Derivative integration.""" +from unittest.mock import patch + import pytest +from homeassistant.components import derivative +from homeassistant.components.derivative.config_flow import ConfigFlowHandler from homeassistant.components.derivative.const import DOMAIN -from homeassistant.core import HomeAssistant +from homeassistant.config_entries import ConfigEntry +from homeassistant.core import Event, HomeAssistant from homeassistant.helpers import device_registry as dr, entity_registry as er +from homeassistant.helpers.event import async_track_entity_registry_updated_event from tests.common import MockConfigEntry -@pytest.mark.parametrize("platform", ["sensor"]) +@pytest.fixture +def sensor_config_entry(hass: HomeAssistant) -> er.RegistryEntry: + """Fixture to create a sensor config entry.""" + sensor_config_entry = MockConfigEntry() + sensor_config_entry.add_to_hass(hass) + return sensor_config_entry + + +@pytest.fixture +def sensor_device( + device_registry: dr.DeviceRegistry, sensor_config_entry: ConfigEntry +) -> dr.DeviceEntry: + """Fixture to create a sensor device.""" + return device_registry.async_get_or_create( + config_entry_id=sensor_config_entry.entry_id, + connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")}, + ) + + +@pytest.fixture +def sensor_entity_entry( + entity_registry: er.EntityRegistry, + sensor_config_entry: ConfigEntry, + sensor_device: dr.DeviceEntry, +) -> er.RegistryEntry: + """Fixture to create a sensor entity entry.""" + return entity_registry.async_get_or_create( + "sensor", + "test", + "unique", + config_entry=sensor_config_entry, + device_id=sensor_device.id, + original_name="ABC", + ) + + +@pytest.fixture +def derivative_config_entry( + hass: HomeAssistant, + sensor_entity_entry: er.RegistryEntry, +) -> MockConfigEntry: + """Fixture to create a derivative config entry.""" + config_entry = MockConfigEntry( + data={}, + domain=DOMAIN, + options={ + "name": "My derivative", + "round": 1.0, + "source": sensor_entity_entry.entity_id, + "time_window": {"seconds": 0.0}, + "unit_prefix": "k", + "unit_time": "min", + }, + title="My derivative", + version=ConfigFlowHandler.VERSION, + minor_version=ConfigFlowHandler.MINOR_VERSION, + ) + + config_entry.add_to_hass(hass) + + return config_entry + + +def track_entity_registry_actions(hass: HomeAssistant, entity_id: str) -> list[str]: + """Track entity registry actions for an entity.""" + events = [] + + def add_event(event: Event[er.EventEntityRegistryUpdatedData]) -> None: + """Add entity registry updated event to the list.""" + events.append(event.data["action"]) + + async_track_entity_registry_updated_event(hass, entity_id, add_event) + + return events + + async def test_setup_and_remove_config_entry( hass: HomeAssistant, entity_registry: er.EntityRegistry, - platform: str, ) -> None: """Test setting up and removing a config entry.""" input_sensor_entity_id = "sensor.input" - derivative_entity_id = f"{platform}.my_derivative" + derivative_entity_id = "sensor.my_derivative" # Setup the config entry config_entry = MockConfigEntry( @@ -147,3 +227,194 @@ async def test_device_cleaning( derivative_config_entry.entry_id ) assert len(devices_after_reload) == 1 + + +async def test_async_handle_source_entity_changes_source_entity_removed( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, + derivative_config_entry: MockConfigEntry, + sensor_config_entry: ConfigEntry, + sensor_device: dr.DeviceEntry, + sensor_entity_entry: er.RegistryEntry, +) -> None: + """Test the derivative config entry is removed when the source entity is removed.""" + # Add another config entry to the sensor device + other_config_entry = MockConfigEntry() + other_config_entry.add_to_hass(hass) + device_registry.async_update_device( + sensor_device.id, add_config_entry_id=other_config_entry.entry_id + ) + + assert await hass.config_entries.async_setup(derivative_config_entry.entry_id) + await hass.async_block_till_done() + + derivative_entity_entry = entity_registry.async_get("sensor.my_derivative") + assert derivative_entity_entry.device_id == sensor_entity_entry.device_id + + sensor_device = device_registry.async_get(sensor_device.id) + assert derivative_config_entry.entry_id in sensor_device.config_entries + + events = track_entity_registry_actions(hass, derivative_entity_entry.entity_id) + + # Remove the source sensor's config entry from the device, this removes the + # source sensor + with patch( + "homeassistant.components.derivative.async_unload_entry", + wraps=derivative.async_unload_entry, + ) as mock_unload_entry: + device_registry.async_update_device( + sensor_device.id, remove_config_entry_id=sensor_config_entry.entry_id + ) + await hass.async_block_till_done() + await hass.async_block_till_done() + mock_unload_entry.assert_called_once() + + # Check that the derivative config entry is removed from the device + sensor_device = device_registry.async_get(sensor_device.id) + assert derivative_config_entry.entry_id not in sensor_device.config_entries + + # Check that the derivative config entry is removed + assert derivative_config_entry.entry_id not in hass.config_entries.async_entry_ids() + + # Check we got the expected events + assert events == ["remove"] + + +async def test_async_handle_source_entity_changes_source_entity_removed_from_device( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, + derivative_config_entry: MockConfigEntry, + sensor_device: dr.DeviceEntry, + sensor_entity_entry: er.RegistryEntry, +) -> None: + """Test the source entity removed from the source device.""" + assert await hass.config_entries.async_setup(derivative_config_entry.entry_id) + await hass.async_block_till_done() + + derivative_entity_entry = entity_registry.async_get("sensor.my_derivative") + assert derivative_entity_entry.device_id == sensor_entity_entry.device_id + + sensor_device = device_registry.async_get(sensor_device.id) + assert derivative_config_entry.entry_id in sensor_device.config_entries + + events = track_entity_registry_actions(hass, derivative_entity_entry.entity_id) + + # Remove the source sensor from the device + with patch( + "homeassistant.components.derivative.async_unload_entry", + wraps=derivative.async_unload_entry, + ) as mock_unload_entry: + entity_registry.async_update_entity( + sensor_entity_entry.entity_id, device_id=None + ) + await hass.async_block_till_done() + mock_unload_entry.assert_called_once() + + # Check that the derivative config entry is removed from the device + sensor_device = device_registry.async_get(sensor_device.id) + assert derivative_config_entry.entry_id not in sensor_device.config_entries + + # Check that the derivative config entry is not removed + assert derivative_config_entry.entry_id in hass.config_entries.async_entry_ids() + + # Check we got the expected events + assert events == ["update"] + + +async def test_async_handle_source_entity_changes_source_entity_moved_other_device( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, + derivative_config_entry: MockConfigEntry, + sensor_config_entry: ConfigEntry, + sensor_device: dr.DeviceEntry, + sensor_entity_entry: er.RegistryEntry, +) -> None: + """Test the source entity is moved to another device.""" + sensor_device_2 = device_registry.async_get_or_create( + config_entry_id=sensor_config_entry.entry_id, + connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:FF")}, + ) + + assert await hass.config_entries.async_setup(derivative_config_entry.entry_id) + await hass.async_block_till_done() + + derivative_entity_entry = entity_registry.async_get("sensor.my_derivative") + assert derivative_entity_entry.device_id == sensor_entity_entry.device_id + + sensor_device = device_registry.async_get(sensor_device.id) + assert derivative_config_entry.entry_id in sensor_device.config_entries + sensor_device_2 = device_registry.async_get(sensor_device_2.id) + assert derivative_config_entry.entry_id not in sensor_device_2.config_entries + + events = track_entity_registry_actions(hass, derivative_entity_entry.entity_id) + + # Move the source sensor to another device + with patch( + "homeassistant.components.derivative.async_unload_entry", + wraps=derivative.async_unload_entry, + ) as mock_unload_entry: + entity_registry.async_update_entity( + sensor_entity_entry.entity_id, device_id=sensor_device_2.id + ) + await hass.async_block_till_done() + mock_unload_entry.assert_called_once() + + # Check that the derivative config entry is moved to the other device + sensor_device = device_registry.async_get(sensor_device.id) + assert derivative_config_entry.entry_id not in sensor_device.config_entries + sensor_device_2 = device_registry.async_get(sensor_device_2.id) + assert derivative_config_entry.entry_id in sensor_device_2.config_entries + + # Check that the derivative config entry is not removed + assert derivative_config_entry.entry_id in hass.config_entries.async_entry_ids() + + # Check we got the expected events + assert events == ["update"] + + +async def test_async_handle_source_entity_new_entity_id( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, + derivative_config_entry: MockConfigEntry, + sensor_device: dr.DeviceEntry, + sensor_entity_entry: er.RegistryEntry, +) -> None: + """Test the source entity's entity ID is changed.""" + assert await hass.config_entries.async_setup(derivative_config_entry.entry_id) + await hass.async_block_till_done() + + derivative_entity_entry = entity_registry.async_get("sensor.my_derivative") + assert derivative_entity_entry.device_id == sensor_entity_entry.device_id + + sensor_device = device_registry.async_get(sensor_device.id) + assert derivative_config_entry.entry_id in sensor_device.config_entries + + events = track_entity_registry_actions(hass, derivative_entity_entry.entity_id) + + # Change the source entity's entity ID + with patch( + "homeassistant.components.derivative.async_unload_entry", + wraps=derivative.async_unload_entry, + ) as mock_unload_entry: + entity_registry.async_update_entity( + sensor_entity_entry.entity_id, new_entity_id="sensor.new_entity_id" + ) + await hass.async_block_till_done() + mock_unload_entry.assert_called_once() + + # Check that the derivative config entry is updated with the new entity ID + assert derivative_config_entry.options["source"] == "sensor.new_entity_id" + + # Check that the helper config is still in the device + sensor_device = device_registry.async_get(sensor_device.id) + assert derivative_config_entry.entry_id in sensor_device.config_entries + + # Check that the derivative config entry is not removed + assert derivative_config_entry.entry_id in hass.config_entries.async_entry_ids() + + # Check we got the expected events + assert events == [] diff --git a/tests/helpers/test_helper_integration.py b/tests/helpers/test_helper_integration.py new file mode 100644 index 00000000000..25d490c27bb --- /dev/null +++ b/tests/helpers/test_helper_integration.py @@ -0,0 +1,424 @@ +"""Tests for the helper entity helpers.""" + +from collections.abc import Generator +from unittest.mock import AsyncMock, Mock + +import pytest + +from homeassistant.config_entries import ConfigEntry +from homeassistant.core import Event, HomeAssistant +from homeassistant.helpers import device_registry as dr, entity_registry as er +from homeassistant.helpers.event import async_track_entity_registry_updated_event +from homeassistant.helpers.helper_integration import async_handle_source_entity_changes + +from tests.common import ( + MockConfigEntry, + MockModule, + mock_config_flow, + mock_integration, + mock_platform, +) + +HELPER_DOMAIN = "helper" +SOURCE_DOMAIN = "test" + + +@pytest.fixture +def source_config_entry(hass: HomeAssistant) -> er.RegistryEntry: + """Fixture to create a source config entry.""" + source_config_entry = MockConfigEntry() + source_config_entry.add_to_hass(hass) + return source_config_entry + + +@pytest.fixture +def source_device( + device_registry: dr.DeviceRegistry, + source_config_entry: ConfigEntry, +) -> dr.DeviceEntry: + """Fixture to create a source device.""" + return device_registry.async_get_or_create( + config_entry_id=source_config_entry.entry_id, + connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")}, + ) + + +@pytest.fixture +def source_entity_entry( + entity_registry: er.EntityRegistry, + source_config_entry: ConfigEntry, + source_device: dr.DeviceEntry, +) -> er.RegistryEntry: + """Fixture to create a source entity entry.""" + return entity_registry.async_get_or_create( + "sensor", + SOURCE_DOMAIN, + "unique", + config_entry=source_config_entry, + device_id=source_device.id, + original_name="ABC", + ) + + +@pytest.fixture +def helper_config_entry( + hass: HomeAssistant, + source_entity_entry: er.RegistryEntry, + use_entity_registry_id: bool, +) -> MockConfigEntry: + """Fixture to create a helper config entry.""" + config_entry = MockConfigEntry( + data={}, + domain=HELPER_DOMAIN, + options={ + "name": "My helper", + "round": 1.0, + "source": source_entity_entry.id + if use_entity_registry_id + else source_entity_entry.entity_id, + "time_window": {"seconds": 0.0}, + "unit_prefix": "k", + "unit_time": "min", + }, + title="My helper", + ) + + config_entry.add_to_hass(hass) + + return config_entry + + +@pytest.fixture +def mock_helper_flow() -> Generator[None]: + """Mock helper config flow.""" + + class MockConfigFlow: + """Mock the helper config flow.""" + + VERSION = 1 + MINOR_VERSION = 1 + + with mock_config_flow(HELPER_DOMAIN, MockConfigFlow): + yield + + +@pytest.fixture +def helper_entity_entry( + entity_registry: er.EntityRegistry, + helper_config_entry: ConfigEntry, + source_device: dr.DeviceEntry, +) -> er.RegistryEntry: + """Fixture to create a helper entity entry.""" + return entity_registry.async_get_or_create( + "sensor", + HELPER_DOMAIN, + helper_config_entry.entry_id, + config_entry=helper_config_entry, + device_id=source_device.id, + original_name="ABC", + ) + + +@pytest.fixture +def async_remove_entry() -> AsyncMock: + """Fixture to mock async_remove_entry.""" + return AsyncMock(return_value=True) + + +@pytest.fixture +def async_unload_entry() -> AsyncMock: + """Fixture to mock async_unload_entry.""" + return AsyncMock(return_value=True) + + +@pytest.fixture +def set_source_entity_id_or_uuid() -> AsyncMock: + """Fixture to mock async_unload_entry.""" + return Mock() + + +@pytest.fixture +def mock_helper_integration( + hass: HomeAssistant, + entity_registry: er.EntityRegistry, + helper_config_entry: MockConfigEntry, + source_entity_entry: er.RegistryEntry, + async_remove_entry: AsyncMock, + async_unload_entry: AsyncMock, + set_source_entity_id_or_uuid: Mock, +) -> 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"], + ) + return True + + mock_integration( + hass, + MockModule( + HELPER_DOMAIN, + async_remove_entry=async_remove_entry, + async_setup_entry=async_setup_entry, + async_unload_entry=async_unload_entry, + ), + ) + mock_platform(hass, f"{HELPER_DOMAIN}.config_flow", None) + + +def track_entity_registry_actions(hass: HomeAssistant, entity_id: str) -> list[str]: + """Track entity registry actions for an entity.""" + events = [] + + def add_event(event: Event[er.EventEntityRegistryUpdatedData]) -> None: + """Add entity registry updated event to the list.""" + events.append(event.data["action"]) + + async_track_entity_registry_updated_event(hass, entity_id, add_event) + + return events + + +@pytest.mark.parametrize("use_entity_registry_id", [True, False]) +@pytest.mark.usefixtures("mock_helper_flow", "mock_helper_integration") +async def test_async_handle_source_entity_changes_source_entity_removed( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, + helper_config_entry: MockConfigEntry, + helper_entity_entry: er.RegistryEntry, + source_config_entry: ConfigEntry, + source_device: dr.DeviceEntry, + source_entity_entry: er.RegistryEntry, + async_remove_entry: AsyncMock, + async_unload_entry: AsyncMock, + set_source_entity_id_or_uuid: Mock, +) -> None: + """Test the helper config entry is removed when the source entity is removed.""" + # Add the helper config entry to the source device + device_registry.async_update_device( + source_device.id, add_config_entry_id=helper_config_entry.entry_id + ) + # Add another config entry to the source device + other_config_entry = MockConfigEntry() + other_config_entry.add_to_hass(hass) + device_registry.async_update_device( + source_device.id, add_config_entry_id=other_config_entry.entry_id + ) + + assert await hass.config_entries.async_setup(helper_config_entry.entry_id) + await hass.async_block_till_done() + + # Check preconditions + helper_entity_entry = entity_registry.async_get(helper_entity_entry.entity_id) + assert helper_entity_entry.device_id == source_entity_entry.device_id + source_device = device_registry.async_get(source_device.id) + assert helper_config_entry.entry_id in source_device.config_entries + + events = track_entity_registry_actions(hass, helper_entity_entry.entity_id) + + # Remove the source entitys's config entry from the device, this removes the + # source entity + device_registry.async_update_device( + source_device.id, remove_config_entry_id=source_config_entry.entry_id + ) + await hass.async_block_till_done() + await hass.async_block_till_done() + + # Check that the helper config entry is unloaded and removed + async_unload_entry.assert_called_once() + async_remove_entry.assert_called_once() + set_source_entity_id_or_uuid.assert_not_called() + + # Check that the helper config entry is removed from the device + source_device = device_registry.async_get(source_device.id) + assert helper_config_entry.entry_id not in source_device.config_entries + + # Check that the helper config entry is removed + assert helper_config_entry.entry_id not in hass.config_entries.async_entry_ids() + + # Check we got the expected events + assert events == ["remove"] + + +@pytest.mark.parametrize("use_entity_registry_id", [True, False]) +@pytest.mark.usefixtures("mock_helper_flow", "mock_helper_integration") +async def test_async_handle_source_entity_changes_source_entity_removed_from_device( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, + helper_config_entry: MockConfigEntry, + helper_entity_entry: er.RegistryEntry, + source_device: dr.DeviceEntry, + source_entity_entry: er.RegistryEntry, + async_remove_entry: AsyncMock, + async_unload_entry: AsyncMock, + set_source_entity_id_or_uuid: Mock, +) -> None: + """Test the source entity removed from the source device.""" + # Add the helper config entry to the source device + device_registry.async_update_device( + source_device.id, add_config_entry_id=helper_config_entry.entry_id + ) + + assert await hass.config_entries.async_setup(helper_config_entry.entry_id) + await hass.async_block_till_done() + + # Check preconditions + helper_entity_entry = entity_registry.async_get(helper_entity_entry.entity_id) + assert helper_entity_entry.device_id == source_entity_entry.device_id + + source_device = device_registry.async_get(source_device.id) + assert helper_config_entry.entry_id in source_device.config_entries + + events = track_entity_registry_actions(hass, helper_entity_entry.entity_id) + + # Remove the source entity from the device + entity_registry.async_update_entity(source_entity_entry.entity_id, device_id=None) + await hass.async_block_till_done() + async_remove_entry.assert_not_called() + async_unload_entry.assert_called_once() + set_source_entity_id_or_uuid.assert_not_called() + + # Check that the helper config entry is removed from the device + source_device = device_registry.async_get(source_device.id) + assert helper_config_entry.entry_id not in source_device.config_entries + + # Check that the helper config entry is not removed + assert helper_config_entry.entry_id in hass.config_entries.async_entry_ids() + + # Check we got the expected events + assert events == ["update"] + + +@pytest.mark.parametrize("use_entity_registry_id", [True, False]) +@pytest.mark.usefixtures("mock_helper_flow", "mock_helper_integration") +async def test_async_handle_source_entity_changes_source_entity_moved_other_device( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, + helper_config_entry: MockConfigEntry, + helper_entity_entry: er.RegistryEntry, + source_config_entry: ConfigEntry, + source_device: dr.DeviceEntry, + source_entity_entry: er.RegistryEntry, + async_remove_entry: AsyncMock, + async_unload_entry: AsyncMock, + set_source_entity_id_or_uuid: Mock, +) -> None: + """Test the source entity is moved to another device.""" + # Add the helper config entry to the source device + device_registry.async_update_device( + source_device.id, add_config_entry_id=helper_config_entry.entry_id + ) + + # Create another device to move the source entity to + source_device_2 = device_registry.async_get_or_create( + config_entry_id=source_config_entry.entry_id, + connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:FF")}, + ) + + assert await hass.config_entries.async_setup(helper_config_entry.entry_id) + await hass.async_block_till_done() + + # Check preconditions + helper_entity_entry = entity_registry.async_get(helper_entity_entry.entity_id) + assert helper_entity_entry.device_id == source_entity_entry.device_id + + source_device = device_registry.async_get(source_device.id) + assert helper_config_entry.entry_id in source_device.config_entries + source_device_2 = device_registry.async_get(source_device_2.id) + assert helper_config_entry.entry_id not in source_device_2.config_entries + + events = track_entity_registry_actions(hass, helper_entity_entry.entity_id) + + # Move the source entity to another device + entity_registry.async_update_entity( + source_entity_entry.entity_id, device_id=source_device_2.id + ) + await hass.async_block_till_done() + async_remove_entry.assert_not_called() + async_unload_entry.assert_called_once() + set_source_entity_id_or_uuid.assert_not_called() + + # Check that the helper config entry is moved to the other device + source_device = device_registry.async_get(source_device.id) + assert helper_config_entry.entry_id not in source_device.config_entries + source_device_2 = device_registry.async_get(source_device_2.id) + assert helper_config_entry.entry_id in source_device_2.config_entries + + # Check that the helper config entry is not removed + assert helper_config_entry.entry_id in hass.config_entries.async_entry_ids() + + # Check we got the expected events + assert events == ["update"] + + +@pytest.mark.parametrize( + ("use_entity_registry_id", "unload_calls", "set_source_entity_id_calls"), + [(True, 1, 0), (False, 0, 1)], +) +@pytest.mark.usefixtures("mock_helper_flow", "mock_helper_integration") +async def test_async_handle_source_entity_new_entity_id( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, + helper_config_entry: MockConfigEntry, + helper_entity_entry: er.RegistryEntry, + source_device: dr.DeviceEntry, + source_entity_entry: er.RegistryEntry, + async_remove_entry: AsyncMock, + async_unload_entry: AsyncMock, + set_source_entity_id_or_uuid: Mock, + unload_calls: int, + set_source_entity_id_calls: int, +) -> None: + """Test the source entity's entity ID is changed.""" + # Add the helper config entry to the source device + device_registry.async_update_device( + source_device.id, add_config_entry_id=helper_config_entry.entry_id + ) + + assert await hass.config_entries.async_setup(helper_config_entry.entry_id) + await hass.async_block_till_done() + + # Check preconditions + helper_entity_entry = entity_registry.async_get(helper_entity_entry.entity_id) + assert helper_entity_entry.device_id == source_entity_entry.device_id + + source_device = device_registry.async_get(source_device.id) + assert helper_config_entry.entry_id in source_device.config_entries + + events = track_entity_registry_actions(hass, helper_entity_entry.entity_id) + + # Change the source entity's entity ID + entity_registry.async_update_entity( + source_entity_entry.entity_id, new_entity_id="sensor.new_entity_id" + ) + await hass.async_block_till_done() + async_remove_entry.assert_not_called() + assert len(async_unload_entry.mock_calls) == unload_calls + assert len(set_source_entity_id_or_uuid.mock_calls) == set_source_entity_id_calls + + # Check that the helper config is still in the device + source_device = device_registry.async_get(source_device.id) + assert helper_config_entry.entry_id in source_device.config_entries + + # Check that the helper config entry is not removed + assert helper_config_entry.entry_id in hass.config_entries.async_entry_ids() + + # Check we got the expected events + assert events == []