Merge pull request #2947 from bramstroker/feat/availability-entity

Add availability entity to control availability of powercalc entity

custom_components/powercalc/config_flow.py

@@ -46,6 +46,7 @@
 from .const import (
     CONF_AREA,
     CONF_AUTOSTART,
+    CONF_AVAILABILITY_ENTITY,
     CONF_CALCULATION_ENABLED_CONDITION,
     CONF_CALIBRATE,
     CONF_CREATE_ENERGY_SENSOR,
@@ -132,7 +133,7 @@
 from .group_include.include import find_entities
 from .power_profile.factory import get_power_profile
 from .power_profile.library import ModelInfo, ProfileLibrary
-from .power_profile.power_profile import DOMAIN_DEVICE_TYPE_MAPPING, SUPPORTED_DOMAINS, DeviceType, PowerProfile
+from .power_profile.power_profile import DEVICE_TYPE_DOMAIN, DOMAIN_DEVICE_TYPE_MAPPING, SUPPORTED_DOMAINS, DeviceType, DiscoveryBy, PowerProfile
 from .sensors.daily_energy import DEFAULT_DAILY_UPDATE_FREQUENCY
 from .sensors.group.config_entry_utils import get_group_entries
 from .sensors.power import PowerSensor
@@ -147,6 +148,7 @@
 class Step(StrEnum):
     ADVANCED_OPTIONS = "advanced_options"
     ASSIGN_GROUPS = "assign_groups"
+    AVAILABILITY_ENTITY = "availability_entity"
     BASIC_OPTIONS = "basic_options"
     GROUP_CUSTOM = "group_custom"
     GROUP_DOMAIN = "group_domain"
@@ -570,7 +572,7 @@ class Step(StrEnum):
 @dataclass(slots=True)
 class PowercalcFormStep:
     schema: vol.Schema | Callable[[], Coroutine[Any, Any, vol.Schema | None]]
-
+    step: Step
     validate_user_input: (
         Callable[
             [dict[str, Any]],
@@ -580,7 +582,6 @@ class PowercalcFormStep:
     ) = None
 
     next_step: Step | Callable[[dict[str, Any]], Coroutine[Any, Any, Step | None]] | None = None
-    step: Step | None = None
     continue_utility_meter_options_step: bool = False
     continue_advanced_step: bool = False
     form_kwarg: dict[str, Any] | None = None
@@ -600,6 +601,7 @@ def __init__(self) -> None:
         self.is_options_flow: bool = isinstance(self, OptionsFlow)
         self.strategy: CalculationStrategy | None = None
         self.name: str | None = None
+        self.handled_steps: list[Step] = []
         super().__init__()
 
     @abstractmethod
@@ -661,42 +663,18 @@ def create_schema_multi_switch(self) -> vol.Schema:
         """Create the config schema for multi switch strategy."""
 
         switch_domains = [str(Platform.SWITCH), str(Platform.LIGHT)]
-        entity_registry = er.async_get(self.hass)
-        entities: list[str] = []
         if self.source_entity and self.source_entity.device_entry:
-            entities = [
-                entity.entity_id
-                for entity in entity_registry.entities.get_entries_for_device_id(self.source_entity.device_entry.id)
-                if entity.domain in switch_domains
-            ]
-
-        # pre-populate the entity selector with the switches from the device
-        if entities:
-            schema = vol.Schema(
-                {
-                    vol.Required(
-                        CONF_ENTITIES,
-                        description={"suggested_value": entities},
-                    ): selector.EntitySelector(
-                        selector.EntitySelectorConfig(
-                            domain=switch_domains,
-                            multiple=True,
-                            include_entities=entities,
-                        ),
-                    ),
-                },
-            )
+            entity_selector = self.create_device_entity_selector(switch_domains, multiple=True)
         else:
-            schema = vol.Schema(
-                {
-                    vol.Required(CONF_ENTITIES): selector.EntitySelector(
-                        selector.EntitySelectorConfig(
-                            domain=switch_domains,
-                            multiple=True,
-                        ),
-                    ),
-                },
+            entity_selector = selector.EntitySelector(
+                selector.EntitySelectorConfig(
+                    domain=switch_domains,
+                    multiple=True,
+                ),
             )
+
+        schema = vol.Schema({vol.Required(CONF_ENTITIES): entity_selector})
+
         if not self.is_library_flow:
             schema = schema.extend(SCHEMA_POWER_MULTI_SWITCH_MANUAL.schema)
 
@@ -797,6 +775,25 @@ def create_source_entity_selector(
             )
         return selector.EntitySelector()
 
+    def create_device_entity_selector(self, domains: list[str], multiple: bool = False) -> selector.EntitySelector:
+        entity_registry = er.async_get(self.hass)
+        if self.source_entity and self.source_entity.device_entry:
+            entities = [
+                entity.entity_id
+                for entity in entity_registry.entities.get_entries_for_device_id(self.source_entity.device_entry.id)
+                if entity.domain in domains
+            ]
+        else:
+            entities = []
+
+        return selector.EntitySelector(
+            selector.EntitySelectorConfig(
+                domain=domains,
+                multiple=multiple,
+                include_entities=entities,
+            ),
+        )
+
     def create_group_selector(
         self,
         current_entry: ConfigEntry | None = None,
@@ -901,6 +898,7 @@ async def handle_form_step(
                 self.name = user_input[CONF_NAME]
             self.sensor_config.update(user_input)
 
+            self.handled_steps.append(form_step.step)
             next_step = form_step.next_step
             if callable(form_step.next_step):
                 next_step = await form_step.next_step(user_input)
@@ -1041,26 +1039,64 @@ async def async_step_post_library(
         if not self.selected_profile:
             return self.async_abort(reason="model_not_supported")  # pragma: no cover
 
-        if self.selected_profile.has_custom_fields and not self.sensor_config.get(CONF_VARIABLES):
+        if Step.LIBRARY_CUSTOM_FIELDS not in self.handled_steps and self.selected_profile.has_custom_fields:
             return await self.async_step_library_custom_fields()
 
-        if await self.selected_profile.has_sub_profiles and not self.selected_profile.sub_profile_select:
+        if Step.AVAILABILITY_ENTITY not in self.handled_steps and self.selected_profile.discovery_by == DiscoveryBy.DEVICE:
+            result = await self.async_step_availability_entity()
+            if result:
+                return result
+
+        if (
+            Step.SUB_PROFILE not in self.handled_steps
+            and await self.selected_profile.has_sub_profiles
+            and not self.selected_profile.sub_profile_select
+        ):
             return await self.async_step_sub_profile()
 
-        if self.selected_profile.device_type == DeviceType.SMART_SWITCH and self.selected_profile.calculation_strategy == CalculationStrategy.FIXED:
+        if (
+            Step.SMART_SWITCH not in self.handled_steps
+            and self.selected_profile.device_type == DeviceType.SMART_SWITCH
+            and self.selected_profile.calculation_strategy == CalculationStrategy.FIXED
+        ):
             return await self.async_step_smart_switch()
 
-        if self.selected_profile.needs_fixed_config:  # pragma: no cover
+        if Step.FIXED not in self.handled_steps and self.selected_profile.needs_fixed_config:  # pragma: no cover
             return await self.async_step_fixed()
 
-        if self.selected_profile.needs_linear_config:
+        if Step.LINEAR not in self.handled_steps and self.selected_profile.needs_linear_config:
             return await self.async_step_linear()
 
-        if self.selected_profile.calculation_strategy == CalculationStrategy.MULTI_SWITCH:
+        if Step.MULTI_SWITCH not in self.handled_steps and self.selected_profile.calculation_strategy == CalculationStrategy.MULTI_SWITCH:
             return await self.async_step_multi_switch()
 
         return await self.async_step_assign_groups()
 
+    async def async_step_availability_entity(self, user_input: dict[str, Any] | None = None) -> FlowResult | None:
+        """Handle the flow for availability entity."""
+        domains = DEVICE_TYPE_DOMAIN[self.selected_profile.device_type]  # type: ignore
+        entity_selector = self.create_device_entity_selector(
+            list(domains) if isinstance(domains, set) else [domains],
+        )
+        try:
+            first_entity = entity_selector.config["include_entities"][0]
+        except IndexError:
+            # Skip step if no entities are available
+            self.handled_steps.append(Step.AVAILABILITY_ENTITY)
+            return None
+        return await self.handle_form_step(
+            PowercalcFormStep(
+                step=Step.AVAILABILITY_ENTITY,
+                schema=vol.Schema(
+                    {
+                        vol.Optional(CONF_AVAILABILITY_ENTITY, default=first_entity): entity_selector,
+                    },
+                ),
+                next_step=Step.POST_LIBRARY,
+            ),
+            user_input,
+        )
+
     async def async_step_library_custom_fields(self, user_input: dict[str, Any] | None = None) -> FlowResult:
         """Handle the flow for custom fields."""
 

custom_components/powercalc/const.py

@@ -39,6 +39,7 @@
 CONF_ALL = "all"
 CONF_AREA = "area"
 CONF_AUTOSTART = "autostart"
+CONF_AVAILABILITY_ENTITY = "availability_entity"
 CONF_CALIBRATE = "calibrate"
 CONF_COMPOSITE = "composite"
 CONF_CREATE_GROUP = "create_group"

custom_components/powercalc/sensor.py

@@ -46,6 +46,7 @@
 )
 from .const import (
     CONF_AND,
+    CONF_AVAILABILITY_ENTITY,
     CONF_CALCULATION_ENABLED_CONDITION,
     CONF_CALIBRATE,
     CONF_COMPOSITE,
@@ -161,6 +162,7 @@
 SENSOR_CONFIG = {
     vol.Optional(CONF_NAME): cv.string,
     vol.Optional(CONF_ENTITY_ID): cv.entity_id,
+    vol.Optional(CONF_AVAILABILITY_ENTITY): cv.entity_id,
     vol.Optional(CONF_UNIQUE_ID): cv.string,
     vol.Optional(CONF_MODEL): cv.string,
     vol.Optional(CONF_MANUFACTURER): cv.string,

custom_components/powercalc/sensors/power.py

@@ -51,6 +51,7 @@
     ATTR_SOURCE_DOMAIN,
     ATTR_SOURCE_ENTITY,
     CALCULATION_STRATEGY_CONF_KEYS,
+    CONF_AVAILABILITY_ENTITY,
     CONF_CALCULATION_ENABLED_CONDITION,
     CONF_CUSTOM_MODEL_DIRECTORY,
     CONF_DELAY,
@@ -375,6 +376,7 @@ def __init__(
         self._standby_sensors: dict = hass.data[DOMAIN][DATA_STANDBY_POWER_SENSORS]
         self.calculation_strategy_factory = calculation_strategy_factory
         self._strategy_instance: PowerCalculationStrategyInterface | None = None
+        self._availability_entity: str | None = sensor_config.get(CONF_AVAILABILITY_ENTITY)
         self._config_entry = config_entry
 
     async def validate(self) -> None:
@@ -484,6 +486,9 @@ def _get_tracking_entities(self) -> list[str | TrackTemplate]:
         if not source_entity_included and self._source_entity.entity_id != DUMMY_ENTITY_ID:
             entities_to_track.append(self._source_entity.entity_id)
 
+        if self._availability_entity and self._availability_entity not in entities_to_track:
+            entities_to_track.append(self._availability_entity)
+
         return entities_to_track
 
     def init_calculation_enabled_condition(self) -> None:
@@ -675,6 +680,10 @@ def native_value(self) -> StateType:
     @property
     def available(self) -> bool:
         """Return True if entity is available."""
+        if self._availability_entity:
+            state = self.hass.states.get(self._availability_entity)
+            return bool(state and state.state != STATE_UNAVAILABLE)
+
         return self._power is not None
 
     def set_energy_sensor_attribute(self, entity_id: str) -> None:

custom_components/powercalc/translations/en.json

@@ -33,6 +33,16 @@
         },
         "title": "Assign to group"
       },
+      "availability_entity": {
+        "data": {
+          "availability_entity": "Availability entity"
+        },
+        "data_description": {
+          "availability_entity": "When this entity is unavailable, the powercalc sensor will be unavailable as well"
+        },
+        "title": "Availability entity",
+        "description": "This profile is per device. Please select the entity which will be used to determine the availability of the powercalc sensor"
+      },
       "daily_energy": {
         "data": {
           "create_utility_meters": "Create utility meters",

tests/config_flow/common.py

@@ -91,6 +91,18 @@ async def select_manufacturer_and_model(
     )
 
 
+async def confirm_auto_discovered_model(
+    hass: HomeAssistant,
+    prev_result: FlowResult,
+    confirmed: bool = True,
+) -> FlowResult:
+    assert prev_result["step_id"] == Step.LIBRARY
+    return await hass.config_entries.flow.async_configure(
+        prev_result["flow_id"],
+        {CONF_CONFIRM_AUTODISCOVERED_MODEL: confirmed},
+    )
+
+
 async def initialize_options_flow(
     hass: HomeAssistant,
     entry: config_entries.ConfigEntry,
@@ -155,11 +167,7 @@ async def initialize_discovery_flow(
     if not confirm_autodiscovered_model:
         return result
 
-    assert result["type"] == data_entry_flow.FlowResultType.FORM
-    return await hass.config_entries.flow.async_configure(
-        result["flow_id"],
-        {CONF_CONFIRM_AUTODISCOVERED_MODEL: True},
-    )
+    return await confirm_auto_discovered_model(hass, result)
 
 
 async def goto_virtual_power_strategy_step(