Fix source device when source entity is changed for Utility Meter (#95636)

* Fix source device when source entity is changed

* Update loop

* Complement and add comments in the test_change_device_source test

* Only clean up dev reg when options change

---------

Co-authored-by: Paulus Schoutsen <balloob@gmail.com>

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."""

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