diff --git a/homeassistant/components/threshold/__init__.py b/homeassistant/components/threshold/__init__.py index fb9e7145951..ea8b469fd32 100644 --- a/homeassistant/components/threshold/__init__.py +++ b/homeassistant/components/threshold/__init__.py @@ -1,13 +1,22 @@ """The threshold component.""" from homeassistant.config_entries import ConfigEntry -from homeassistant.const import Platform +from homeassistant.const import CONF_ENTITY_ID, Platform from homeassistant.core import HomeAssistant -from homeassistant.helpers import device_registry as dr +from homeassistant.helpers.device import ( + async_remove_stale_devices_links_keep_entity_device, +) async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up Min/Max from a config entry.""" + + async_remove_stale_devices_links_keep_entity_device( + hass, + entry.entry_id, + entry.options[CONF_ENTITY_ID], + ) + await hass.config_entries.async_forward_entry_setups( entry, (Platform.BINARY_SENSOR,) ) @@ -20,16 +29,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def config_entry_update_listener(hass: HomeAssistant, entry: ConfigEntry) -> None: """Update listener, called when the config entry options are changed.""" - # Remove device link for entry, the source device may have changed. - # The link will be recreated after load. - device_registry = dr.async_get(hass) - devices = device_registry.devices.get_devices_for_config_entry_id(entry.entry_id) - - for device in devices: - device_registry.async_update_device( - device.id, remove_config_entry_id=entry.entry_id - ) - await hass.config_entries.async_reload(entry.entry_id) diff --git a/homeassistant/components/threshold/binary_sensor.py b/homeassistant/components/threshold/binary_sensor.py index ac970a53f55..8c3882ff360 100644 --- a/homeassistant/components/threshold/binary_sensor.py +++ b/homeassistant/components/threshold/binary_sensor.py @@ -30,11 +30,8 @@ from homeassistant.core import ( HomeAssistant, callback, ) -from homeassistant.helpers import ( - config_validation as cv, - device_registry as dr, - entity_registry as er, -) +from homeassistant.helpers import config_validation as cv, entity_registry as er +from homeassistant.helpers.device import async_device_info_to_link_from_entity from homeassistant.helpers.device_registry import DeviceInfo from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.event import async_track_state_change_event @@ -87,27 +84,10 @@ async def async_setup_entry( registry, config_entry.options[CONF_ENTITY_ID] ) - source_entity = registry.async_get(entity_id) - dev_reg = dr.async_get(hass) - # Resolve source entity device - if ( - (source_entity is not None) - and (source_entity.device_id is not None) - and ( - ( - device := dev_reg.async_get( - device_id=source_entity.device_id, - ) - ) - is not None - ) - ): - device_info = DeviceInfo( - identifiers=device.identifiers, - connections=device.connections, - ) - else: - device_info = None + device_info = async_device_info_to_link_from_entity( + hass, + entity_id, + ) hysteresis = config_entry.options[CONF_HYSTERESIS] lower = config_entry.options[CONF_LOWER] diff --git a/tests/components/threshold/test_init.py b/tests/components/threshold/test_init.py index d1fda706911..6e85d659922 100644 --- a/tests/components/threshold/test_init.py +++ b/tests/components/threshold/test_init.py @@ -122,3 +122,89 @@ async def test_entry_changed(hass: HomeAssistant, platform) -> None: # Check that the config entry association has updated assert config_entry.entry_id not in _get_device_config_entries(run1_entry) assert config_entry.entry_id in _get_device_config_entries(run2_entry) + + +async def test_device_cleaning( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, +) -> None: + """Test for source entity device for Threshold.""" + + # Source entity device config entry + source_config_entry = MockConfigEntry() + source_config_entry.add_to_hass(hass) + + # Device entry of the source entity + source_device1_entry = device_registry.async_get_or_create( + config_entry_id=source_config_entry.entry_id, + identifiers={("sensor", "identifier_test1")}, + connections={("mac", "30:31:32:33:34:01")}, + ) + + # Source entity registry + source_entity = entity_registry.async_get_or_create( + "sensor", + "test", + "source", + config_entry=source_config_entry, + device_id=source_device1_entry.id, + ) + await hass.async_block_till_done() + assert entity_registry.async_get("sensor.test_source") is not None + + # Configure the configuration entry for Threshold + threshold_config_entry = MockConfigEntry( + data={}, + domain=DOMAIN, + options={ + "entity_id": "sensor.test_source", + "hysteresis": 0.0, + "lower": -2.0, + "name": "Threshold", + "upper": None, + }, + title="Threshold", + ) + threshold_config_entry.add_to_hass(hass) + assert await hass.config_entries.async_setup(threshold_config_entry.entry_id) + await hass.async_block_till_done() + + # Confirm the link between the source entity device and the threshold sensor + threshold_entity = entity_registry.async_get("binary_sensor.threshold") + assert threshold_entity is not None + assert threshold_entity.device_id == source_entity.device_id + + # Device entry incorrectly linked to Threshold config entry + device_registry.async_get_or_create( + config_entry_id=threshold_config_entry.entry_id, + identifiers={("sensor", "identifier_test2")}, + connections={("mac", "30:31:32:33:34:02")}, + ) + device_registry.async_get_or_create( + config_entry_id=threshold_config_entry.entry_id, + identifiers={("sensor", "identifier_test3")}, + connections={("mac", "30:31:32:33:34:03")}, + ) + await hass.async_block_till_done() + + # Before reloading the config entry, two devices are expected to be linked + devices_before_reload = device_registry.devices.get_devices_for_config_entry_id( + threshold_config_entry.entry_id + ) + assert len(devices_before_reload) == 3 + + # Config entry reload + await hass.config_entries.async_reload(threshold_config_entry.entry_id) + await hass.async_block_till_done() + + # Confirm the link between the source entity device and the threshold sensor after reload + threshold_entity = entity_registry.async_get("binary_sensor.threshold") + assert threshold_entity is not None + assert threshold_entity.device_id == source_entity.device_id + + # After reloading the config entry, only one linked device is expected + devices_after_reload = device_registry.devices.get_devices_for_config_entry_id( + threshold_config_entry.entry_id + ) + assert len(devices_after_reload) == 1