Merge aa507e4 into 2aa067b

custom_components/powercalc/config_flow.py

@@ -721,6 +721,8 @@ async def save_options(self, user_input: dict[str, Any]) -> dict:
                 if user_input and key in user_input:
                     generic_options[key] = user_input.get(key)
 
+            if CONF_ENTITY_ID in user_input:
+                generic_options[CONF_ENTITY_ID] = user_input[CONF_ENTITY_ID]
             self.current_config.update(generic_options)
 
             if self.strategy:
@@ -763,9 +765,20 @@ def build_options_schema(self) -> vol.Schema:
                 if self.strategy
                 else vol.Schema({})
             )
-            data_schema = SCHEMA_POWER_OPTIONS.extend(strategy_schema.schema).extend(
-                SCHEMA_POWER_ADVANCED.schema,
+
+            data_schema = (
+                vol.Schema(
+                    {
+                        vol.Optional(CONF_ENTITY_ID): _create_source_entity_selector(
+                            False,
+                        ),
+                    },
+                )
+                .extend(SCHEMA_POWER_OPTIONS.schema)
+                .extend(strategy_schema.schema)
+                .extend(SCHEMA_POWER_ADVANCED.schema)
             )
+
             strategy_options = self.current_config.get(self.strategy) or {}
 
         if self.sensor_type == SensorType.DAILY_ENERGY:
@@ -808,20 +821,23 @@ def _get_strategy_schema(strategy: str, source_entity_id: str) -> vol.Schema:
     return vol.Schema({})
 
 
-def _create_virtual_power_schema(
-    hass: HomeAssistant,
-    is_library_flow: bool = True,
-) -> vol.Schema:
+def _create_source_entity_selector(is_library_flow: bool = True) -> selector.EntitySelector:
     if is_library_flow:
-        entity_selector = selector.EntitySelector(
+        return selector.EntitySelector(
             selector.EntitySelectorConfig(domain=list(DEVICE_DOMAINS.values())),
         )
-    else:
-        entity_selector = selector.EntitySelector()
+    return selector.EntitySelector()
 
+
+def _create_virtual_power_schema(
+    hass: HomeAssistant,
+    is_library_flow: bool = True,
+) -> vol.Schema:
     schema = vol.Schema(
         {
-            vol.Required(CONF_ENTITY_ID): entity_selector,
+            vol.Required(CONF_ENTITY_ID): _create_source_entity_selector(
+                is_library_flow,
+            ),
         },
     ).extend(SCHEMA_POWER_BASE.schema)
     schema = schema.extend({vol.Optional(CONF_GROUP): _create_group_selector(hass)})

custom_components/powercalc/sensors/abstract.py

@@ -35,7 +35,9 @@ async def async_added_to_hass(self) -> None:
         device_id: str = getattr(self, "source_device_id")  # noqa: B009
         device_reg = dr.async_get(self.hass)
         device_entry = device_reg.async_get(device_id)
-        if not device_entry or device_entry.id == entity_entry.device_id:  # pragma: no cover
+        if (
+            not device_entry or device_entry.id == entity_entry.device_id
+        ):  # pragma: no cover
             return
         _LOGGER.debug(f"Binding {self.entity_id} to device {device_id}")
         entity_reg.async_update_entity(self.entity_id, device_id=device_id)

tests/test_sensor.py

@@ -787,3 +787,77 @@ async def test_rename_source_entity_id(hass: HomeAssistant) -> None:
     power_state = hass.states.get("sensor.testentry_power")
     assert power_state.state == "50.00"
     assert power_state.attributes.get(ATTR_SOURCE_ENTITY) == new_light_id
+
+
+async def test_change_config_entry_entity_id(hass: HomeAssistant) -> None:
+    """Test that changing the source entity of an existing config entry works correctly"""
+
+    original_light_id = "light.original"
+    original_unique_id = "aaaa"
+    new_light_id = "light.new"
+    new_unique_id = "bbbb"
+    mock_registry(
+        hass,
+        {
+            original_light_id: er.RegistryEntry(
+                entity_id=original_light_id,
+                unique_id=original_unique_id,
+                platform="light",
+            ),
+            new_light_id: er.RegistryEntry(
+                entity_id=original_light_id,
+                unique_id=new_unique_id,
+                platform="light",
+            ),
+        },
+    )
+
+    # Create an existing config entry referencing the original source entity
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        data={
+            CONF_SENSOR_TYPE: SensorType.VIRTUAL_POWER,
+            CONF_CREATE_ENERGY_SENSOR: True,
+            CONF_CREATE_UTILITY_METERS: True,
+            CONF_MODE: CalculationStrategy.FIXED,
+            CONF_NAME: "testentry",
+            CONF_ENTITY_ID: original_light_id,
+            CONF_FIXED: {CONF_POWER: 50},
+            CONF_UNIQUE_ID: original_unique_id,
+        },
+        unique_id=original_unique_id,
+    )
+    entry.add_to_hass(hass)
+
+    await run_powercalc_setup(
+        hass,
+        {},
+    )
+
+    hass.states.async_set(original_light_id, STATE_ON)
+    hass.states.async_set(new_light_id, STATE_ON)
+    await hass.async_block_till_done()
+
+    power_state = hass.states.get("sensor.testentry_power")
+    assert power_state.attributes.get(ATTR_SOURCE_ENTITY) == original_light_id
+    assert power_state.state == "50.00"
+
+    # Change the entity_id using the options flow
+    result = await hass.config_entries.options.async_init(
+        entry.entry_id,
+        data=None,
+    )
+    user_input = {CONF_ENTITY_ID: new_light_id, CONF_POWER: 100}
+    await hass.config_entries.options.async_configure(
+        result["flow_id"],
+        user_input=user_input,
+    )
+    await hass.async_block_till_done()
+
+    changed_entry = hass.config_entries.async_get_entry(entry.entry_id)
+    assert changed_entry.data.get(CONF_ENTITY_ID) == new_light_id
+    assert changed_entry.unique_id == original_unique_id
+
+    power_state = hass.states.get("sensor.testentry_power")
+    assert power_state.attributes.get(ATTR_SOURCE_ENTITY) == new_light_id
+    assert power_state.state == "100.00"