Merge fcfcdf0 into cbe2c55

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,
@@ -263,8 +264,24 @@ async def async_setup_entry(
         )
         async_add_entities(entities)
         return
+
+    # Add entry to an existing group
+    reload_group_entry = None
+    if sensor_type == SensorType.VIRTUAL_POWER and CONF_GROUP in sensor_config:
+        group_entry_id = sensor_config.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 not entry.entry_id in member_sensors:
+            member_sensors.append(entry.entry_id)
+            hass.config_entries.async_update_entry(
+                group_entry,
+                data={**group_entry.data, CONF_GROUP_MEMBER_SENSORS: member_sensors}
+            )
+            reload_group_entry = group_entry
 
     await _async_setup_entities(hass, sensor_config, async_add_entities)
+    if reload_group_entry:
+        await hass.config_entries.async_reload(group_entry.entry_id)
 
 
 async def _async_setup_entities(

custom_components/powercalc/sensors/group.py

@@ -159,7 +159,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 = (

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,
@@ -500,3 +501,49 @@ 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_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_group = MockConfigEntry(
+        domain=DOMAIN,
+        data={
+            CONF_SENSOR_TYPE: SensorType.GROUP,
+            CONF_NAME: "GroupA",
+            CONF_GROUP_POWER_ENTITIES: ["sensor.test1_power"],
+            CONF_GROUP_ENERGY_ENTITIES: ["sensor.test1_energy"],
+        },
+    )
+    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_sensor = MockConfigEntry(
+        domain=DOMAIN,
+        data={
+            CONF_SENSOR_TYPE: SensorType.VIRTUAL_POWER,
+            CONF_NAME: "VirtualSensor",
+            CONF_ENTITY_ID: DUMMY_ENTITY_ID,
+            CONF_UNIQUE_ID: "abcdef",
+            CONF_GROUP: config_entry_group.entry_id,
+            CONF_FIXED: {
+                CONF_POWER: 50
+            }
+        },
+    )
+    config_entry_sensor.add_to_hass(hass)
+    assert await hass.config_entries.async_setup(config_entry_sensor.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_sensor.entry_id]
+
+    group_state = hass.states.get("sensor.groupa_power")
+    assert group_state
+    assert group_state.attributes.get("entities") == {
+        "sensor.test1_power",
+        "sensor.virtualsensor_power"
+    }