From 296b5c627a441ca29ff9aa055e04dbe5692f7d22 Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Wed, 11 Jun 2025 15:18:04 +0200 Subject: [PATCH] Handle changes to source entities in generic_hygrostat helper (#146538) --- .../components/generic_hygrostat/__init__.py | 59 ++- .../components/generic_hygrostat/test_init.py | 422 +++++++++++++++++- 2 files changed, 478 insertions(+), 3 deletions(-) diff --git a/homeassistant/components/generic_hygrostat/__init__.py b/homeassistant/components/generic_hygrostat/__init__.py index b4a6014c5a4..a12994c1a75 100644 --- a/homeassistant/components/generic_hygrostat/__init__.py +++ b/homeassistant/components/generic_hygrostat/__init__.py @@ -5,11 +5,18 @@ import voluptuous as vol from homeassistant.components.humidifier import HumidifierDeviceClass from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_NAME, CONF_UNIQUE_ID, Platform -from homeassistant.core import HomeAssistant -from homeassistant.helpers import config_validation as cv, discovery +from homeassistant.core import Event, HomeAssistant +from homeassistant.helpers import ( + config_validation as cv, + discovery, + 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.event import async_track_entity_registry_updated_event +from homeassistant.helpers.helper_integration import async_handle_source_entity_changes from homeassistant.helpers.typing import ConfigType DOMAIN = "generic_hygrostat" @@ -88,6 +95,54 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: entry.options[CONF_HUMIDIFIER], ) + def set_humidifier_entity_id_or_uuid(source_entity_id: str) -> None: + hass.config_entries.async_update_entry( + entry, + options={**entry.options, CONF_HUMIDIFIER: 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( + # We use async_handle_source_entity_changes to track changes to the humidifer, + # but not the humidity sensor because the generic_hygrostat adds itself to the + # humidifier's device. + async_handle_source_entity_changes( + hass, + helper_config_entry_id=entry.entry_id, + set_source_entity_id_or_uuid=set_humidifier_entity_id_or_uuid, + source_device_id=async_entity_id_to_device_id( + hass, entry.options[CONF_HUMIDIFIER] + ), + source_entity_id_or_uuid=entry.options[CONF_HUMIDIFIER], + source_entity_removed=source_entity_removed, + ) + ) + + async def async_sensor_updated( + event: Event[er.EventEntityRegistryUpdatedData], + ) -> None: + """Handle entity registry update.""" + data = event.data + if data["action"] != "update": + return + if "entity_id" not in data["changes"]: + return + + # Entity_id changed, update the config entry + hass.config_entries.async_update_entry( + entry, + options={**entry.options, CONF_SENSOR: data["entity_id"]}, + ) + + entry.async_on_unload( + async_track_entity_registry_updated_event( + hass, entry.options[CONF_SENSOR], async_sensor_updated + ) + ) + await hass.config_entries.async_forward_entry_setups(entry, (Platform.HUMIDIFIER,)) entry.async_on_unload(entry.add_update_listener(config_entry_update_listener)) return True diff --git a/tests/components/generic_hygrostat/test_init.py b/tests/components/generic_hygrostat/test_init.py index 16bb4dc6db5..254d4da5806 100644 --- a/tests/components/generic_hygrostat/test_init.py +++ b/tests/components/generic_hygrostat/test_init.py @@ -2,15 +2,136 @@ from __future__ import annotations +from unittest.mock import patch + +import pytest + +from homeassistant.components import generic_hygrostat from homeassistant.components.generic_hygrostat import DOMAIN -from homeassistant.core import HomeAssistant +from homeassistant.components.generic_hygrostat.config_flow import ConfigFlowHandler +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 .test_humidifier import ENT_SENSOR 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:EE")}, + ) + + +@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 switch_config_entry(hass: HomeAssistant) -> er.RegistryEntry: + """Fixture to create a switch config entry.""" + switch_config_entry = MockConfigEntry() + switch_config_entry.add_to_hass(hass) + return switch_config_entry + + +@pytest.fixture +def switch_device( + device_registry: dr.DeviceRegistry, switch_config_entry: ConfigEntry +) -> dr.DeviceEntry: + """Fixture to create a switch device.""" + return device_registry.async_get_or_create( + config_entry_id=switch_config_entry.entry_id, + connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")}, + ) + + +@pytest.fixture +def switch_entity_entry( + entity_registry: er.EntityRegistry, + switch_config_entry: ConfigEntry, + switch_device: dr.DeviceEntry, +) -> er.RegistryEntry: + """Fixture to create a switch entity entry.""" + return entity_registry.async_get_or_create( + "switch", + "test", + "unique", + config_entry=switch_config_entry, + device_id=switch_device.id, + original_name="ABC", + ) + + +@pytest.fixture +def generic_hygrostat_config_entry( + hass: HomeAssistant, + sensor_entity_entry: er.RegistryEntry, + switch_entity_entry: er.RegistryEntry, +) -> MockConfigEntry: + """Fixture to create a generic_hygrostat config entry.""" + config_entry = MockConfigEntry( + data={}, + domain=DOMAIN, + options={ + "device_class": "humidifier", + "dry_tolerance": 2.0, + "humidifier": switch_entity_entry.entity_id, + "name": "My generic hygrostat", + "target_sensor": sensor_entity_entry.entity_id, + "wet_tolerance": 4.0, + }, + title="My generic hygrostat", + 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_device_cleaning( hass: HomeAssistant, device_registry: dr.DeviceRegistry, @@ -98,3 +219,302 @@ async def test_device_cleaning( assert len(devices_after_reload) == 1 assert devices_after_reload[0].id == source_device1_entry.id + + +@pytest.mark.usefixtures( + "sensor_config_entry", + "sensor_device", + "sensor_entity_entry", + "switch_config_entry", + "switch_device", +) +@pytest.mark.parametrize( + ("source_entity_id", "helper_in_device", "expected_events"), + [("switch.test_unique", True, ["update"]), ("sensor.test_unique", False, [])], +) +async def test_async_handle_source_entity_changes_source_entity_removed( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, + generic_hygrostat_config_entry: MockConfigEntry, + switch_entity_entry: er.RegistryEntry, + source_entity_id: str, + helper_in_device: bool, + expected_events: list[str], +) -> None: + """Test the generic_hygrostat config entry is removed when the source entity is removed.""" + source_entity_entry = entity_registry.async_get(source_entity_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_entity_entry.device_id, add_config_entry_id=other_config_entry.entry_id + ) + + assert await hass.config_entries.async_setup( + generic_hygrostat_config_entry.entry_id + ) + await hass.async_block_till_done() + + generic_hygrostat_entity_entry = entity_registry.async_get( + "humidifier.my_generic_hygrostat" + ) + assert generic_hygrostat_entity_entry.device_id == switch_entity_entry.device_id + + source_device = device_registry.async_get(source_entity_entry.device_id) + assert ( + generic_hygrostat_config_entry.entry_id in source_device.config_entries + ) == helper_in_device + + events = track_entity_registry_actions( + hass, generic_hygrostat_entity_entry.entity_id + ) + + # Remove the source entity's config entry from the device, this removes the + # source entity + with patch( + "homeassistant.components.generic_hygrostat.async_unload_entry", + wraps=generic_hygrostat.async_unload_entry, + ) as mock_unload_entry: + device_registry.async_update_device( + source_device.id, remove_config_entry_id=source_entity_entry.config_entry_id + ) + await hass.async_block_till_done() + await hass.async_block_till_done() + mock_unload_entry.assert_not_called() + + # Check if the generic_hygrostat config entry is not in the device + source_device = device_registry.async_get(source_device.id) + assert generic_hygrostat_config_entry.entry_id not in source_device.config_entries + + # Check that the generic_hygrostat config entry is not removed + assert ( + generic_hygrostat_config_entry.entry_id in hass.config_entries.async_entry_ids() + ) + + # Check we got the expected events + assert events == expected_events + + +@pytest.mark.usefixtures( + "sensor_config_entry", + "sensor_device", + "sensor_entity_entry", + "switch_config_entry", + "switch_device", +) +@pytest.mark.parametrize( + ("source_entity_id", "helper_in_device", "unload_entry_calls", "expected_events"), + [("switch.test_unique", True, 1, ["update"]), ("sensor.test_unique", False, 0, [])], +) +async def test_async_handle_source_entity_changes_source_entity_removed_from_device( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, + generic_hygrostat_config_entry: MockConfigEntry, + switch_entity_entry: er.RegistryEntry, + source_entity_id: str, + helper_in_device: bool, + unload_entry_calls: int, + expected_events: list[str], +) -> None: + """Test the source entity removed from the source device.""" + source_entity_entry = entity_registry.async_get(source_entity_id) + + assert await hass.config_entries.async_setup( + generic_hygrostat_config_entry.entry_id + ) + await hass.async_block_till_done() + + generic_hygrostat_entity_entry = entity_registry.async_get( + "humidifier.my_generic_hygrostat" + ) + assert generic_hygrostat_entity_entry.device_id == switch_entity_entry.device_id + + source_device = device_registry.async_get(source_entity_entry.device_id) + assert ( + generic_hygrostat_config_entry.entry_id in source_device.config_entries + ) == helper_in_device + + events = track_entity_registry_actions( + hass, generic_hygrostat_entity_entry.entity_id + ) + + # Remove the source entity from the device + with patch( + "homeassistant.components.generic_hygrostat.async_unload_entry", + wraps=generic_hygrostat.async_unload_entry, + ) as mock_unload_entry: + entity_registry.async_update_entity( + source_entity_entry.entity_id, device_id=None + ) + await hass.async_block_till_done() + assert len(mock_unload_entry.mock_calls) == unload_entry_calls + + # Check that the generic_hygrostat config entry is removed from the device + source_device = device_registry.async_get(source_device.id) + assert generic_hygrostat_config_entry.entry_id not in source_device.config_entries + + # Check that the generic_hygrostat config entry is not removed + assert ( + generic_hygrostat_config_entry.entry_id in hass.config_entries.async_entry_ids() + ) + + # Check we got the expected events + assert events == expected_events + + +@pytest.mark.usefixtures( + "sensor_config_entry", + "sensor_device", + "sensor_entity_entry", + "switch_config_entry", + "switch_device", +) +@pytest.mark.parametrize( + ("source_entity_id", "helper_in_device", "unload_entry_calls", "expected_events"), + [("switch.test_unique", True, 1, ["update"]), ("sensor.test_unique", False, 0, [])], +) +async def test_async_handle_source_entity_changes_source_entity_moved_other_device( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, + generic_hygrostat_config_entry: MockConfigEntry, + switch_entity_entry: er.RegistryEntry, + source_entity_id: str, + helper_in_device: bool, + unload_entry_calls: int, + expected_events: list[str], +) -> None: + """Test the source entity is moved to another device.""" + source_entity_entry = entity_registry.async_get(source_entity_id) + + source_device_2 = device_registry.async_get_or_create( + config_entry_id=source_entity_entry.config_entry_id, + connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:FF")}, + ) + + assert await hass.config_entries.async_setup( + generic_hygrostat_config_entry.entry_id + ) + await hass.async_block_till_done() + + generic_hygrostat_entity_entry = entity_registry.async_get( + "humidifier.my_generic_hygrostat" + ) + assert generic_hygrostat_entity_entry.device_id == switch_entity_entry.device_id + + source_device = device_registry.async_get(source_entity_entry.device_id) + assert ( + generic_hygrostat_config_entry.entry_id in source_device.config_entries + ) == helper_in_device + source_device_2 = device_registry.async_get(source_device_2.id) + assert generic_hygrostat_config_entry.entry_id not in source_device_2.config_entries + + events = track_entity_registry_actions( + hass, generic_hygrostat_entity_entry.entity_id + ) + + # Move the source entity to another device + with patch( + "homeassistant.components.generic_hygrostat.async_unload_entry", + wraps=generic_hygrostat.async_unload_entry, + ) as mock_unload_entry: + entity_registry.async_update_entity( + source_entity_entry.entity_id, device_id=source_device_2.id + ) + await hass.async_block_till_done() + assert len(mock_unload_entry.mock_calls) == unload_entry_calls + + # Check that the generic_hygrostat config entry is moved to the other device + source_device = device_registry.async_get(source_device.id) + assert generic_hygrostat_config_entry.entry_id not in source_device.config_entries + source_device_2 = device_registry.async_get(source_device_2.id) + assert ( + generic_hygrostat_config_entry.entry_id in source_device_2.config_entries + ) == helper_in_device + + # Check that the generic_hygrostat config entry is not removed + assert ( + generic_hygrostat_config_entry.entry_id in hass.config_entries.async_entry_ids() + ) + + # Check we got the expected events + assert events == expected_events + + +@pytest.mark.usefixtures( + "sensor_config_entry", + "sensor_device", + "sensor_entity_entry", + "switch_config_entry", + "switch_device", +) +@pytest.mark.parametrize( + ("source_entity_id", "new_entity_id", "helper_in_device", "config_key"), + [ + ("switch.test_unique", "switch.new_entity_id", True, "humidifier"), + ("sensor.test_unique", "sensor.new_entity_id", False, "target_sensor"), + ], +) +async def test_async_handle_source_entity_new_entity_id( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, + generic_hygrostat_config_entry: MockConfigEntry, + switch_entity_entry: er.RegistryEntry, + source_entity_id: str, + new_entity_id: str, + helper_in_device: bool, + config_key: str, +) -> None: + """Test the source entity's entity ID is changed.""" + source_entity_entry = entity_registry.async_get(source_entity_id) + + assert await hass.config_entries.async_setup( + generic_hygrostat_config_entry.entry_id + ) + await hass.async_block_till_done() + + generic_hygrostat_entity_entry = entity_registry.async_get( + "humidifier.my_generic_hygrostat" + ) + assert generic_hygrostat_entity_entry.device_id == switch_entity_entry.device_id + + source_device = device_registry.async_get(source_entity_entry.device_id) + assert ( + generic_hygrostat_config_entry.entry_id in source_device.config_entries + ) == helper_in_device + + events = track_entity_registry_actions( + hass, generic_hygrostat_entity_entry.entity_id + ) + + # Change the source entity's entity ID + with patch( + "homeassistant.components.generic_hygrostat.async_unload_entry", + wraps=generic_hygrostat.async_unload_entry, + ) as mock_unload_entry: + entity_registry.async_update_entity( + source_entity_entry.entity_id, new_entity_id=new_entity_id + ) + await hass.async_block_till_done() + mock_unload_entry.assert_called_once() + + # Check that the generic_hygrostat config entry is updated with the new entity ID + assert generic_hygrostat_config_entry.options[config_key] == new_entity_id + + # Check that the helper config is still in the device + source_device = device_registry.async_get(source_device.id) + assert ( + generic_hygrostat_config_entry.entry_id in source_device.config_entries + ) == helper_in_device + + # Check that the generic_hygrostat config entry is not removed + assert ( + generic_hygrostat_config_entry.entry_id in hass.config_entries.async_entry_ids() + ) + + # Check we got the expected events + assert events == []