Option in config flow to add virtual power sensor to group on creation

custom_components/powercalc/__init__.py

@@ -73,7 +73,7 @@
     get_power_profile,
     has_manufacturer_and_model_information,
 )
-from .sensors.group import create_group_sensors
+from .sensors.group import create_group_sensors, update_associated_group_entry
 from .strategy.factory import PowerCalculatorStrategyFactory
 
 PLATFORMS = [Platform.SENSOR]
@@ -220,6 +220,10 @@ async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
     )
 
     if unload_ok:
+        updated_group_entry = await update_associated_group_entry(hass, config_entry, remove=True)
+        if updated_group_entry:
+            await hass.config_entries.async_reload(updated_group_entry.entry_id)
+
         used_unique_ids: list[str] = hass.data[DOMAIN][DATA_USED_UNIQUE_IDS]
         try:
             used_unique_ids.remove(config_entry.unique_id)

custom_components/powercalc/config_flow.py

@@ -1,6 +1,7 @@
 """Config flow for Adaptive Lighting integration."""
 
 from __future__ import annotations
+from audioop import mul
 
 import copy
 import logging
@@ -21,7 +22,7 @@
     Platform,
 )
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.data_entry_flow import FlowResult
+from homeassistant.data_entry_flow import FlowResult, FlowHandler
 from homeassistant.helpers import selector
 
 from .common import SourceEntity, create_source_entity
@@ -32,6 +33,7 @@
     CONF_DAILY_FIXED_ENERGY,
     CONF_FIXED,
     CONF_GAMMA_CURVE,
+    CONF_GROUP,
     CONF_GROUP_ENERGY_ENTITIES,
     CONF_GROUP_MEMBER_SENSORS,
     CONF_GROUP_POWER_ENTITIES,
@@ -116,7 +118,7 @@
     }
 )
 
-SCHEMA_POWER = vol.Schema(
+SCHEMA_POWER_BASE = vol.Schema(
     {
         vol.Required(CONF_ENTITY_ID): selector.EntitySelector(),
         vol.Optional(CONF_NAME): selector.TextSelector(),
@@ -135,7 +137,7 @@
             )
         ),
     }
-).extend(SCHEMA_POWER_OPTIONS.schema)
+)
 
 SCHEMA_POWER_FIXED = vol.Schema(
     {
@@ -225,7 +227,7 @@ async def async_step_virtual_power(
 
         return self.async_show_form(
             step_id="virtual_power",
-            data_schema=SCHEMA_POWER,
+            data_schema=_create_virtual_power_schema(self.hass),
             errors={},
         )
 
@@ -440,7 +442,6 @@ def create_config_entry(self) -> FlowResult:
             self.sensor_config.update({CONF_UNIQUE_ID: self.unique_id})
         return self.async_create_entry(title=self.name, data=self.sensor_config)
 
-
 class OptionsFlowHandler(OptionsFlow):
     """Handle an option flow for PowerCalc."""
 
@@ -525,10 +526,9 @@ def build_options_schema(self) -> vol.Schema:
 
         strategy_options = {}
         if self.sensor_type == SensorType.VIRTUAL_POWER:
-            base_power_schema = SCHEMA_POWER_OPTIONS
             strategy: str = self.current_config.get(CONF_MODE)
             strategy_schema = _get_strategy_schema(strategy, self.source_entity_id)
-            data_schema = base_power_schema.extend(strategy_schema.schema)
+            data_schema = SCHEMA_POWER_OPTIONS.extend(strategy_schema.schema)
             strategy_options = self.current_config.get(strategy) or {}
 
         if self.sensor_type == SensorType.DAILY_ENERGY:
@@ -568,23 +568,16 @@ def _get_strategy_schema(strategy: str, source_entity_id: str) -> vol.Schema:
     if strategy == CalculationStrategy.LUT:
         return vol.Schema({})
 
+def _create_virtual_power_schema(hass: HomeAssistant) -> vol.Schema:
+    base_schema: vol.Schema = SCHEMA_POWER_BASE.extend(
+        {
+            vol.Optional(CONF_GROUP): _create_group_selector(hass)
+        }
+    )
+    return base_schema.extend(SCHEMA_POWER_OPTIONS.schema)
 
 def _create_group_options_schema(hass: HomeAssistant) -> vol.Schema:
     """Create config schema for groups"""
-    sub_groups = [
-        selector.SelectOptionDict(
-            value=config_entry.entry_id, label=config_entry.data.get(CONF_NAME)
-        )
-        for config_entry in hass.config_entries.async_entries(DOMAIN)
-        if config_entry.data.get(CONF_SENSOR_TYPE) == SensorType.GROUP
-    ]
-
-    sub_group_selector = selector.SelectSelector(
-        selector.SelectSelectorConfig(
-            options=sub_groups, multiple=True, mode=selector.SelectSelectorMode.DROPDOWN
-        )
-    )
-
     member_sensors = [
         selector.SelectOptionDict(
             value=config_entry.entry_id, label=config_entry.data.get(CONF_NAME)
@@ -618,14 +611,29 @@ def _create_group_options_schema(hass: HomeAssistant) -> vol.Schema:
                     multiple=True,
                 )
             ),
-            vol.Optional(CONF_SUB_GROUPS): sub_group_selector,
+            vol.Optional(CONF_SUB_GROUPS): _create_group_selector(hass, multiple=True),
             vol.Optional(
                 CONF_CREATE_UTILITY_METERS, default=False
             ): selector.BooleanSelector(),
             vol.Optional(CONF_HIDE_MEMBERS, default=False): selector.BooleanSelector(),
         }
     )
 
+def _create_group_selector(hass: HomeAssistant, multiple: bool = False) -> selector.SelectSelector:
+    options = [
+        selector.SelectOptionDict(
+            value=config_entry.entry_id, label=config_entry.data.get(CONF_NAME)
+        )
+        for config_entry in hass.config_entries.async_entries(DOMAIN)
+        if config_entry.data.get(CONF_SENSOR_TYPE) == SensorType.GROUP
+    ]
+
+    return selector.SelectSelector(
+        selector.SelectSelectorConfig(
+            options=options, multiple=multiple, mode=selector.SelectSelectorMode.DROPDOWN
+        )
+    )
+
 
 def _validate_group_input(user_input: dict[str, str] = None) -> dict:
     """Validate the group form"""

custom_components/powercalc/sensor.py

@@ -79,6 +79,7 @@
     CONF_ENERGY_SENSOR_UNIT_PREFIX,
     CONF_FIXED,
     CONF_GROUP,
+    CONF_GROUP_MEMBER_SENSORS,
     CONF_HIDE_MEMBERS,
     CONF_IGNORE_UNAVAILABLE_STATE,
     CONF_INCLUDE,
@@ -130,7 +131,7 @@
     create_daily_fixed_energy_sensor,
 )
 from .sensors.energy import create_energy_sensor
-from .sensors.group import create_group_sensors, create_group_sensors_from_config_entry
+from .sensors.group import create_group_sensors, create_group_sensors_from_config_entry, update_associated_group_entry
 from .sensors.power import RealPowerSensor, VirtualPowerSensor, create_power_sensor
 from .sensors.utility_meter import create_utility_meters
 from .strategy.fixed import CONFIG_SCHEMA as FIXED_SCHEMA
@@ -263,8 +264,13 @@ async def async_setup_entry(
         )
         async_add_entities(entities)
         return
+
+    # Add entry to an existing group
+    updated_group_entry = await update_associated_group_entry(hass, entry, remove=False)
 
     await _async_setup_entities(hass, sensor_config, async_add_entities)
+    if updated_group_entry:
+        await hass.config_entries.async_reload(updated_group_entry.entry_id)
 
 
 async def _async_setup_entities(

custom_components/powercalc/sensors/group.py

@@ -35,14 +35,17 @@
     ATTR_IS_GROUP,
     CONF_ENERGY_SENSOR_PRECISION,
     CONF_ENERGY_SENSOR_UNIT_PREFIX,
+    CONF_GROUP,
     CONF_GROUP_ENERGY_ENTITIES,
     CONF_GROUP_MEMBER_SENSORS,
     CONF_GROUP_POWER_ENTITIES,
     CONF_HIDE_MEMBERS,
     CONF_POWER_SENSOR_PRECISION,
+    CONF_SENSOR_TYPE,
     CONF_SUB_GROUPS,
     DOMAIN,
     SERVICE_RESET_ENERGY,
+    SensorType,
     UnitPrefix,
 )
 from .abstract import (
@@ -142,6 +145,33 @@ async def create_group_sensors_from_config_entry(
 
     return group_sensors
 
+async def update_associated_group_entry(hass: HomeAssistant, config_entry: ConfigEntry, remove: bool) -> ConfigEntry | None:
+    """
+    Update the group config entry when the virtual power config entry is associated to a group
+    Adds the sensor to the group on creation of the config entry
+    Removes the sensor from the group on removal of the config entry
+    """
+    sensor_type = config_entry.data.get(CONF_SENSOR_TYPE)
+    if sensor_type != SensorType.VIRTUAL_POWER:
+        return None
+    if CONF_GROUP not in config_entry.data:
+        return None
+
+    group_entry_id = config_entry.data.get(CONF_GROUP)
+    group_entry = hass.config_entries.async_get_entry(group_entry_id)
+    member_sensors = group_entry.data.get(CONF_GROUP_MEMBER_SENSORS) or []
+
+    if remove and config_entry.entry_id in member_sensors:
+        member_sensors.remove(config_entry.entry_id)
+    elif config_entry.entry_id not in member_sensors:
+        member_sensors.append(config_entry.entry_id)
+
+    hass.config_entries.async_update_entry(
+        group_entry,
+        data={**group_entry.data, CONF_GROUP_MEMBER_SENSORS: member_sensors}
+    )
+    return group_entry
+
 
 @callback
 def resolve_entity_ids_recursively(
@@ -159,7 +189,7 @@ def resolve_entity_ids_recursively(
 
     # Include the power/energy sensors for an existing Virtual Power config entry
     entity_reg = er.async_get(hass)
-    member_entry_ids = entry.data.get(CONF_GROUP_MEMBER_SENSORS)
+    member_entry_ids = entry.data.get(CONF_GROUP_MEMBER_SENSORS) or []
     # Unfortunately device_class is not correctly set at this time in the entity_registry
     # So we need to match on state_class.
     state_class = (

custom_components/powercalc/strings.json

@@ -54,6 +54,7 @@
           "mode": "Calculation strategy",
           "entity_id": "Source entity",
           "unique_id": "Unique id",
+          "group": "Add to group",
           "standby_power": "Standby power",
           "create_energy_sensor": "Create energy sensor",
           "create_utility_meters": "Create utility meters"

custom_components/powercalc/translations/en.json

@@ -121,6 +121,7 @@
                     "create_energy_sensor": "Create energy sensor",
                     "create_utility_meters": "Create utility meters",
                     "entity_id": "Source entity",
+                    "group": "Add to group",
                     "mode": "Calculation strategy",
                     "name": "Name",
                     "standby_power": "Standby power",

custom_components/powercalc/translations/nl.json

@@ -121,6 +121,7 @@
                     "create_energy_sensor": "Creëer energie sensoren",
                     "create_utility_meters": "Creëer utiliteit meters",
                     "entity_id": "Bron entiteit",
+                    "group": "Aan groep toevoegen",
                     "mode": "Calculatie strategie",
                     "name": "Naam",
                     "standby_power": "Standby stroom",

tests/sensors/test_group.py

@@ -36,6 +36,7 @@
     CONF_CREATE_UTILITY_METERS,
     CONF_ENERGY_SENSOR_UNIT_PREFIX,
     CONF_FIXED,
+    CONF_GROUP,
     CONF_GROUP_ENERGY_ENTITIES,
     CONF_GROUP_MEMBER_SENSORS,
     CONF_GROUP_POWER_ENTITIES,
@@ -56,6 +57,7 @@
 from ..common import (
     create_input_boolean,
     create_input_booleans,
+    create_mocked_virtual_power_sensor_entry,
     get_simple_fixed_config,
     run_powercalc_setup_yaml_config,
 )
@@ -500,3 +502,69 @@ async def test_include_config_entries_in_group(hass: HomeAssistant):
     assert group_energy_state.attributes.get(ATTR_ENTITIES) == {
         "sensor.virtualsensor_energy"
     }
+
+async def test_add_virtual_power_sensor_to_group_on_creation(hass: HomeAssistant):
+    """
+    When creating a virtual power sensor using the config flow you can define a group you want to add it to
+    Test that the new sensors are added to the existing group correctly
+    """
+
+    config_entry_sensor1 = await create_mocked_virtual_power_sensor_entry(hass, "VirtualSensor1", "xyz")
+
+    config_entry_group = MockConfigEntry(
+        domain=DOMAIN,
+        data={
+            CONF_SENSOR_TYPE: SensorType.GROUP,
+            CONF_NAME: "GroupA",
+            CONF_GROUP_MEMBER_SENSORS: [config_entry_sensor1.entry_id]
+        },
+    )
+    config_entry_group.add_to_hass(hass)
+    assert await hass.config_entries.async_setup(config_entry_group.entry_id)
+    await hass.async_block_till_done()
+
+    config_entry_sensor2 = MockConfigEntry(
+        domain=DOMAIN,
+        data={
+            CONF_SENSOR_TYPE: SensorType.VIRTUAL_POWER,
+            CONF_NAME: "VirtualSensor2",
+            CONF_ENTITY_ID: DUMMY_ENTITY_ID,
+            CONF_UNIQUE_ID: "abc",
+            CONF_GROUP: config_entry_group.entry_id,
+            CONF_FIXED: {
+                CONF_POWER: 50
+            }
+        },
+    )
+    config_entry_sensor2.add_to_hass(hass)
+    assert await hass.config_entries.async_setup(config_entry_sensor2.entry_id)
+    await hass.async_block_till_done()
+
+    config_entry_group = hass.config_entries.async_get_entry(config_entry_group.entry_id)
+    assert config_entry_group.data.get(CONF_GROUP_MEMBER_SENSORS) == [
+        config_entry_sensor1.entry_id,
+        config_entry_sensor2.entry_id
+    ]
+
+    group_state = hass.states.get("sensor.groupa_power")
+    assert group_state
+    assert group_state.attributes.get("entities") == {
+        "sensor.virtualsensor1_power",
+        "sensor.virtualsensor2_power"
+    }
+
+    # Remove config entry from Home Assistant, and see if group is updated accordingly
+    await hass.config_entries.async_remove(config_entry_sensor2.entry_id)
+    await hass.async_block_till_done()
+
+    config_entry_group = hass.config_entries.async_get_entry(config_entry_group.entry_id)
+    assert config_entry_group.data.get(CONF_GROUP_MEMBER_SENSORS) == [
+        config_entry_sensor1.entry_id,
+    ]
+
+    group_state = hass.states.get("sensor.groupa_power")
+    assert group_state
+    assert group_state.attributes.get("entities") == {
+        "sensor.virtualsensor1_power",
+    }
+