diff --git a/homeassistant/components/utility_meter/__init__.py b/homeassistant/components/utility_meter/__init__.py index 11e58fca775..ffe6d7f5433 100644 --- a/homeassistant/components/utility_meter/__init__.py +++ b/homeassistant/components/utility_meter/__init__.py @@ -10,7 +10,11 @@ from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN from homeassistant.config_entries import ConfigEntry from homeassistant.const import ATTR_ENTITY_ID, CONF_NAME, CONF_UNIQUE_ID, Platform from homeassistant.core import HomeAssistant, split_entity_id -from homeassistant.helpers import discovery, entity_registry as er +from homeassistant.helpers import ( + device_registry as dr, + discovery, + entity_registry as er, +) import homeassistant.helpers.config_validation as cv from homeassistant.helpers.dispatcher import async_dispatcher_send from homeassistant.helpers.typing import ConfigType @@ -182,7 +186,9 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up Utility Meter from a config entry.""" entity_registry = er.async_get(hass) - hass.data[DATA_UTILITY][entry.entry_id] = {} + hass.data[DATA_UTILITY][entry.entry_id] = { + "source": entry.options[CONF_SOURCE_SENSOR], + } hass.data[DATA_UTILITY][entry.entry_id][DATA_TARIFF_SENSORS] = [] try: @@ -218,8 +224,23 @@ 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.""" + old_source = hass.data[DATA_UTILITY][entry.entry_id]["source"] await hass.config_entries.async_reload(entry.entry_id) + if old_source == entry.options[CONF_SOURCE_SENSOR]: + return + + entity_registry = er.async_get(hass) + device_registry = dr.async_get(hass) + + old_source_entity = entity_registry.async_get(old_source) + if not old_source_entity or not old_source_entity.device_id: + return + + device_registry.async_update_device( + old_source_entity.device_id, remove_config_entry_id=entry.entry_id + ) + async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Unload a config entry.""" diff --git a/tests/components/utility_meter/test_config_flow.py b/tests/components/utility_meter/test_config_flow.py index 302d3879a04..88a77407c07 100644 --- a/tests/components/utility_meter/test_config_flow.py +++ b/tests/components/utility_meter/test_config_flow.py @@ -7,6 +7,7 @@ from homeassistant import config_entries from homeassistant.components.utility_meter.const import DOMAIN from homeassistant.core import HomeAssistant from homeassistant.data_entry_flow import FlowResultType +from homeassistant.helpers import device_registry as dr, entity_registry as er from tests.common import MockConfigEntry @@ -266,3 +267,173 @@ async def test_options(hass: HomeAssistant) -> None: await hass.async_block_till_done() state = hass.states.get("sensor.electricity_meter") assert state.attributes["source"] == input_sensor2_entity_id + + +async def test_change_device_source(hass: HomeAssistant) -> None: + """Test remove the device registry configuration entry when the source entity changes.""" + + device_registry = dr.async_get(hass) + entity_registry = er.async_get(hass) + + # Configure source entity 1 (with a linked device) + source_config_entry_1 = MockConfigEntry() + source_device_entry_1 = device_registry.async_get_or_create( + config_entry_id=source_config_entry_1.entry_id, + identifiers={("sensor", "identifier_test1")}, + connections={("mac", "20:31:32:33:34:35")}, + ) + source_entity_1 = entity_registry.async_get_or_create( + "sensor", + "test", + "source1", + config_entry=source_config_entry_1, + device_id=source_device_entry_1.id, + ) + + # Configure source entity 2 (with a linked device) + source_config_entry_2 = MockConfigEntry() + source_device_entry_2 = device_registry.async_get_or_create( + config_entry_id=source_config_entry_2.entry_id, + identifiers={("sensor", "identifier_test2")}, + connections={("mac", "30:31:32:33:34:35")}, + ) + source_entity_2 = entity_registry.async_get_or_create( + "sensor", + "test", + "source2", + config_entry=source_config_entry_2, + device_id=source_device_entry_2.id, + ) + + # Configure source entity 3 (without a device) + source_config_entry_3 = MockConfigEntry() + source_entity_3 = entity_registry.async_get_or_create( + "sensor", + "test", + "source3", + config_entry=source_config_entry_3, + ) + + await hass.async_block_till_done() + + input_sensor_entity_id_1 = "sensor.test_source1" + input_sensor_entity_id_2 = "sensor.test_source2" + input_sensor_entity_id_3 = "sensor.test_source3" + + # Test the existence of configured source entities + assert entity_registry.async_get(input_sensor_entity_id_1) is not None + assert entity_registry.async_get(input_sensor_entity_id_2) is not None + assert entity_registry.async_get(input_sensor_entity_id_3) is not None + + # Setup the config entry with source entity 1 (with a linked device) + current_entity_source = source_entity_1 + utility_meter_config_entry = MockConfigEntry( + data={}, + domain=DOMAIN, + options={ + "cycle": "monthly", + "delta_values": False, + "name": "Energy", + "net_consumption": False, + "offset": 0, + "periodically_resetting": True, + "source": current_entity_source.entity_id, + "tariffs": [], + }, + title="Energy", + ) + utility_meter_config_entry.add_to_hass(hass) + assert await hass.config_entries.async_setup(utility_meter_config_entry.entry_id) + await hass.async_block_till_done() + + # Confirm that the configuration entry has been added to the source entity 1 (current) device registry + current_device = device_registry.async_get( + device_id=current_entity_source.device_id + ) + assert utility_meter_config_entry.entry_id in current_device.config_entries + + # Change configuration options to use source entity 2 (with a linked device) and reload the integration + previous_entity_source = source_entity_1 + current_entity_source = source_entity_2 + + result = await hass.config_entries.options.async_init( + utility_meter_config_entry.entry_id + ) + assert result["type"] == FlowResultType.FORM + result = await hass.config_entries.options.async_configure( + result["flow_id"], + user_input={ + "periodically_resetting": True, + "source": current_entity_source.entity_id, + }, + ) + assert result["type"] == FlowResultType.CREATE_ENTRY + await hass.async_block_till_done() + + # Confirm that the configuration entry has been removed from the source entity 1 (previous) device registry + previous_device = device_registry.async_get( + device_id=previous_entity_source.device_id + ) + assert utility_meter_config_entry.entry_id not in previous_device.config_entries + + # Confirm that the configuration entry has been added to the source entity 2 (current) device registry + current_device = device_registry.async_get( + device_id=current_entity_source.device_id + ) + assert utility_meter_config_entry.entry_id in current_device.config_entries + + # Change configuration options to use source entity 3 (without a device) and reload the integration + previous_entity_source = source_entity_2 + current_entity_source = source_entity_3 + + result = await hass.config_entries.options.async_init( + utility_meter_config_entry.entry_id + ) + assert result["type"] == FlowResultType.FORM + result = await hass.config_entries.options.async_configure( + result["flow_id"], + user_input={ + "periodically_resetting": True, + "source": current_entity_source.entity_id, + }, + ) + assert result["type"] == FlowResultType.CREATE_ENTRY + await hass.async_block_till_done() + + # Confirm that the configuration entry has been removed from the source entity 2 (previous) device registry + previous_device = device_registry.async_get( + device_id=previous_entity_source.device_id + ) + assert utility_meter_config_entry.entry_id not in previous_device.config_entries + + # Confirm that there is no device with the helper configuration entry + assert ( + dr.async_entries_for_config_entry( + device_registry, utility_meter_config_entry.entry_id + ) + == [] + ) + + # Change configuration options to use source entity 2 (with a linked device) and reload the integration + previous_entity_source = source_entity_3 + current_entity_source = source_entity_2 + + result = await hass.config_entries.options.async_init( + utility_meter_config_entry.entry_id + ) + assert result["type"] == FlowResultType.FORM + result = await hass.config_entries.options.async_configure( + result["flow_id"], + user_input={ + "periodically_resetting": True, + "source": current_entity_source.entity_id, + }, + ) + assert result["type"] == FlowResultType.CREATE_ENTRY + await hass.async_block_till_done() + + # Confirm that the configuration entry has been added to the source entity 2 (current) device registry + current_device = device_registry.async_get( + device_id=current_entity_source.device_id + ) + assert utility_meter_config_entry.entry_id in current_device.config_entries