diff --git a/homeassistant/components/threshold/__init__.py b/homeassistant/components/threshold/__init__.py index ea8b469fd32..9460a50db80 100644 --- a/homeassistant/components/threshold/__init__.py +++ b/homeassistant/components/threshold/__init__.py @@ -4,8 +4,10 @@ from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_ENTITY_ID, Platform from homeassistant.core import HomeAssistant 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 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: @@ -17,6 +19,29 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: entry.options[CONF_ENTITY_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}, + ) + + 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) + + entry.async_on_unload( + async_handle_source_entity_changes( + hass, + helper_config_entry_id=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_ENTITY_ID] + ), + source_entity_id_or_uuid=entry.options[CONF_ENTITY_ID], + source_entity_removed=source_entity_removed, + ) + ) + await hass.config_entries.async_forward_entry_setups( entry, (Platform.BINARY_SENSOR,) ) diff --git a/tests/components/threshold/test_init.py b/tests/components/threshold/test_init.py index 6e85d659922..599612ce0b7 100644 --- a/tests/components/threshold/test_init.py +++ b/tests/components/threshold/test_init.py @@ -1,14 +1,95 @@ """Test the Min/Max integration.""" +from unittest.mock import patch + import pytest +from homeassistant.components import threshold +from homeassistant.components.threshold.config_flow import ConfigFlowHandler from homeassistant.components.threshold.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.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 threshold_config_entry( + hass: HomeAssistant, + sensor_entity_entry: er.RegistryEntry, +) -> MockConfigEntry: + """Fixture to create a threshold config entry.""" + config_entry = MockConfigEntry( + data={}, + domain=DOMAIN, + options={ + "entity_id": sensor_entity_entry.entity_id, + "hysteresis": 0.0, + "lower": -2.0, + "name": "My threshold", + "upper": None, + }, + title="My threshold", + 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 + + @pytest.mark.parametrize("platform", ["binary_sensor"]) async def test_setup_and_remove_config_entry( hass: HomeAssistant, @@ -208,3 +289,194 @@ async def test_device_cleaning( threshold_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, + threshold_config_entry: MockConfigEntry, + sensor_config_entry: ConfigEntry, + sensor_device: dr.DeviceEntry, + sensor_entity_entry: er.RegistryEntry, +) -> None: + """Test the threshold 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(threshold_config_entry.entry_id) + await hass.async_block_till_done() + + threshold_entity_entry = entity_registry.async_get("binary_sensor.my_threshold") + assert threshold_entity_entry.device_id == sensor_entity_entry.device_id + + sensor_device = device_registry.async_get(sensor_device.id) + assert threshold_config_entry.entry_id in sensor_device.config_entries + + events = track_entity_registry_actions(hass, threshold_entity_entry.entity_id) + + # Remove the source sensor's config entry from the device, this removes the + # source sensor + with patch( + "homeassistant.components.threshold.async_unload_entry", + wraps=threshold.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_not_called() + + # Check that the threshold config entry is removed from the device + sensor_device = device_registry.async_get(sensor_device.id) + assert threshold_config_entry.entry_id not in sensor_device.config_entries + + # Check that the threshold config entry is not removed + assert threshold_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_removed_from_device( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, + threshold_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(threshold_config_entry.entry_id) + await hass.async_block_till_done() + + threshold_entity_entry = entity_registry.async_get("binary_sensor.my_threshold") + assert threshold_entity_entry.device_id == sensor_entity_entry.device_id + + sensor_device = device_registry.async_get(sensor_device.id) + assert threshold_config_entry.entry_id in sensor_device.config_entries + + events = track_entity_registry_actions(hass, threshold_entity_entry.entity_id) + + # Remove the source sensor from the device + with patch( + "homeassistant.components.threshold.async_unload_entry", + wraps=threshold.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 threshold config entry is removed from the device + sensor_device = device_registry.async_get(sensor_device.id) + assert threshold_config_entry.entry_id not in sensor_device.config_entries + + # Check that the threshold config entry is not removed + assert threshold_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, + threshold_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(threshold_config_entry.entry_id) + await hass.async_block_till_done() + + threshold_entity_entry = entity_registry.async_get("binary_sensor.my_threshold") + assert threshold_entity_entry.device_id == sensor_entity_entry.device_id + + sensor_device = device_registry.async_get(sensor_device.id) + assert threshold_config_entry.entry_id in sensor_device.config_entries + sensor_device_2 = device_registry.async_get(sensor_device_2.id) + assert threshold_config_entry.entry_id not in sensor_device_2.config_entries + + events = track_entity_registry_actions(hass, threshold_entity_entry.entity_id) + + # Move the source sensor to another device + with patch( + "homeassistant.components.threshold.async_unload_entry", + wraps=threshold.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 threshold config entry is moved to the other device + sensor_device = device_registry.async_get(sensor_device.id) + assert threshold_config_entry.entry_id not in sensor_device.config_entries + sensor_device_2 = device_registry.async_get(sensor_device_2.id) + assert threshold_config_entry.entry_id in sensor_device_2.config_entries + + # Check that the threshold config entry is not removed + assert threshold_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, + threshold_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(threshold_config_entry.entry_id) + await hass.async_block_till_done() + + threshold_entity_entry = entity_registry.async_get("binary_sensor.my_threshold") + assert threshold_entity_entry.device_id == sensor_entity_entry.device_id + + sensor_device = device_registry.async_get(sensor_device.id) + assert threshold_config_entry.entry_id in sensor_device.config_entries + + events = track_entity_registry_actions(hass, threshold_entity_entry.entity_id) + + # Change the source entity's entity ID + with patch( + "homeassistant.components.threshold.async_unload_entry", + wraps=threshold.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 threshold config entry is updated with the new entity ID + assert threshold_config_entry.options["entity_id"] == "sensor.new_entity_id" + + # Check that the helper config is still in the device + sensor_device = device_registry.async_get(sensor_device.id) + assert threshold_config_entry.entry_id in sensor_device.config_entries + + # Check that the threshold config entry is not removed + assert threshold_config_entry.entry_id in hass.config_entries.async_entry_ids() + + # Check we got the expected events + assert events == []