Allow to use playbook strategy in composite (#2551)

* fix: fully implement interface in composite implementation

* chore: add failing test

* feat: implement playbook support in composite strategy

* feat: implement playbook support in composite strategy

* feat: implement playbook support in composite strategy

* feat: implement playbook support in composite strategy

custom_components/powercalc/sensors/power.py

@@ -462,7 +462,7 @@ async def initial_update(hass: HomeAssistant) -> None:
 
         self.async_on_remove(start.async_at_start(self.hass, initial_update))
 
-        if isinstance(self._strategy_instance, PlaybookStrategy):
+        if hasattr(self._strategy_instance, "set_update_callback"):
             self._strategy_instance.set_update_callback(self._update_power_sensor)
 
         @callback
@@ -616,7 +616,7 @@ def _update_sleep_power(*_: Any) -> None:  # noqa: ANN401
 
         standby_power = self._standby_power
         if self._strategy_instance.can_calculate_standby():
-            standby_power = await self._strategy_instance.calculate(state) or Decimal(0)
+            standby_power = await self._strategy_instance.calculate(state) or self._standby_power
 
         evaluated = await evaluate_power(standby_power)
         if evaluated is None:

custom_components/powercalc/strategy/composite.py

@@ -1,15 +1,17 @@
 from __future__ import annotations
 
 import logging
+from collections.abc import Callable
 from dataclasses import dataclass
 from decimal import Decimal
 
-from homeassistant.const import CONF_ENTITY_ID
+from homeassistant.const import CONF_ENTITY_ID, STATE_OFF
 from homeassistant.core import HomeAssistant, State
 from homeassistant.helpers.condition import ConditionCheckerType
 from homeassistant.helpers.event import TrackTemplate
 from homeassistant.helpers.template import Template
 
+from .playbook import PlaybookStrategy
 from .strategy_interface import PowerCalculationStrategyInterface
 
 _LOGGER = logging.getLogger(__name__)
@@ -19,20 +21,60 @@ class CompositeStrategy(PowerCalculationStrategyInterface):
     def __init__(self, hass: HomeAssistant, strategies: list[SubStrategy]) -> None:
         self.hass = hass
         self.strategies = strategies
+        self.playbook_strategies: list[PlaybookStrategy] = [
+            strategy.strategy for strategy in self.strategies if isinstance(strategy.strategy, PlaybookStrategy)
+        ]
 
     async def calculate(self, entity_state: State) -> Decimal | None:
+        """Calculate power consumption based on entity state."""
+        await self.stop_active_playbooks()
+
         for sub_strategy in self.strategies:
-            if sub_strategy.condition and not sub_strategy.condition(
-                self.hass,
-                {"state": entity_state},
-            ):
+            strategy = sub_strategy.strategy
+
+            if sub_strategy.condition and not sub_strategy.condition(self.hass, {"state": entity_state}):
                 continue
 
-            return await sub_strategy.strategy.calculate(entity_state)
+            if isinstance(strategy, PlaybookStrategy):
+                await self.activate_playbook(strategy)
+
+            if entity_state.state == STATE_OFF and strategy.can_calculate_standby():
+                return await strategy.calculate(entity_state)
+
+            if entity_state.state != STATE_OFF:
+                return await strategy.calculate(entity_state)
 
         return None
 
+    async def stop_active_playbooks(self) -> None:
+        """Stop any active playbooks from sub strategies."""
+        for playbook in self.playbook_strategies:
+            await playbook.stop_playbook()
+
+    @staticmethod
+    async def activate_playbook(strategy: PlaybookStrategy) -> None:
+        """Activate the first playbook in the list."""
+        if not strategy.registered_playbooks:
+            return  # pragma: no cover
+        playbook = strategy.registered_playbooks[0]
+        await strategy.activate_playbook(playbook)
+
+    def set_update_callback(self, update_callback: Callable[[Decimal], None]) -> None:
+        """
+        Register update callback which allows to give the strategy instance access to the power sensor
+        and manipulate the state
+        """
+        for sub_strategy in self.strategies:
+            if hasattr(sub_strategy.strategy, "set_update_callback"):
+                sub_strategy.strategy.set_update_callback(update_callback)
+
+    async def validate_config(self) -> None:
+        """Validate correct setup of the strategy."""
+        for sub_strategy in self.strategies:
+            await sub_strategy.strategy.validate_config()
+
     def get_entities_to_track(self) -> list[str | TrackTemplate]:
+        """Return entities that should be tracked."""
         track_templates: list[str | TrackTemplate] = []
         for sub_strategy in self.strategies:
             if sub_strategy.condition_config:
@@ -42,11 +84,21 @@ def get_entities_to_track(self) -> list[str | TrackTemplate]:
                 )
         return track_templates
 
+    def can_calculate_standby(self) -> bool:
+        """Return if this strategy can calculate standby power."""
+        return any(sub_strategy.strategy.can_calculate_standby() for sub_strategy in self.strategies)
+
+    async def on_start(self, hass: HomeAssistant) -> None:
+        """Called after HA has started"""
+        for sub_strategy in self.strategies:
+            await sub_strategy.strategy.on_start(hass)
+
     def resolve_track_templates_from_condition(
         self,
         condition_config: dict,
         templates: list[str | TrackTemplate],
     ) -> None:
+        """Resolve track templates from condition config."""
         for key, value in condition_config.items():
             if key == CONF_ENTITY_ID and isinstance(value, list):
                 templates.extend(value)

custom_components/powercalc/strategy/multi_switch.py

@@ -45,7 +45,7 @@ async def calculate(self, entity_state: State) -> Decimal | None:
                 entity_id: (state.state if (state := self.hass.states.get(entity_id)) else STATE_UNAVAILABLE) for entity_id in self.switch_entities
             }
 
-        if entity_state.entity_id != DUMMY_ENTITY_ID:
+        if entity_state.entity_id != DUMMY_ENTITY_ID and entity_state.entity_id in self.switch_entities:
             self.known_states[entity_state.entity_id] = entity_state.state
 
         def _get_power(state: str) -> Decimal:

custom_components/powercalc/strategy/playbook.py

@@ -203,6 +203,11 @@ def _load_playbook_entries() -> list[PlaybookEntry]:
     def can_calculate_standby(self) -> bool:
         return bool(self._states_trigger and STATE_OFF in self._states_trigger)
 
+    @property
+    def registered_playbooks(self) -> list[str]:
+        playbooks = dict(self._config.get(CONF_PLAYBOOKS, {}))
+        return list(playbooks.keys())
+
 
 class PlaybookQueue:
     def __init__(self, items: list[PlaybookEntry]) -> None:

tests/strategy/test_composite.py

@@ -1,25 +1,36 @@
+from datetime import timedelta
+
 from homeassistant.components.light import ATTR_BRIGHTNESS
 from homeassistant.const import (
     CONF_CONDITION,
+    CONF_ENTITIES,
     CONF_ENTITY_ID,
+    CONF_NAME,
     STATE_OFF,
     STATE_ON,
     STATE_UNAVAILABLE,
 )
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.device_registry import DeviceEntry
 from homeassistant.helpers.entity_registry import RegistryEntry
-from pytest_homeassistant_custom_component.common import mock_device_registry, mock_registry
+from homeassistant.util import dt
+from pytest_homeassistant_custom_component.common import async_fire_time_changed, mock_device_registry, mock_registry
 
 from custom_components.powercalc.const import (
     CONF_COMPOSITE,
     CONF_FIXED,
     CONF_LINEAR,
     CONF_MAX_POWER,
     CONF_MIN_POWER,
+    CONF_MULTI_SWITCH,
+    CONF_PLAYBOOK,
+    CONF_PLAYBOOKS,
     CONF_POWER,
+    CONF_POWER_OFF,
+    CONF_STANDBY_POWER,
 )
 from tests.common import (
+    get_test_config_dir,
     run_powercalc_setup,
 )
 
@@ -217,3 +228,160 @@ async def test_nested_conditions(hass: HomeAssistant) -> None:
     await hass.async_block_till_done()
 
     assert hass.states.get("sensor.test_power").state == STATE_UNAVAILABLE
+
+
+async def test_playbook(hass: HomeAssistant) -> None:
+    hass.config.config_dir = get_test_config_dir()
+
+    dishwasher_mode_entity = "sensor.dishwasher_operating_mode"
+
+    hass.states.async_set(dishwasher_mode_entity, "Cycle Complete")
+    await hass.async_block_till_done()
+
+    sensor_config = {
+        CONF_ENTITY_ID: dishwasher_mode_entity,
+        CONF_NAME: "Dishwasher",
+        CONF_COMPOSITE: [
+            {
+                CONF_CONDITION: {
+                    "condition": "state",
+                    "entity_id": dishwasher_mode_entity,
+                    "state": "Cycle Active",
+                },
+                CONF_PLAYBOOK: {
+                    CONF_PLAYBOOKS: {
+                        "playbook": "composite/dishwasher.csv",
+                    },
+                },
+            },
+            {
+                CONF_CONDITION: {
+                    "condition": "state",
+                    "entity_id": dishwasher_mode_entity,
+                    "state": "Cycle Complete",
+                },
+                CONF_FIXED: {
+                    CONF_POWER: 9.6,
+                },
+            },
+            {
+                CONF_FIXED: {
+                    CONF_POWER: 1.6,
+                },
+            },
+        ],
+    }
+
+    await run_powercalc_setup(hass, sensor_config, {})
+
+    hass.states.async_set(dishwasher_mode_entity, "Cycle Complete")
+    await hass.async_block_till_done()
+
+    assert hass.states.get("sensor.dishwasher_power").state == "9.60"
+
+    hass.states.async_set(dishwasher_mode_entity, "Cycle Paused")
+    await hass.async_block_till_done()
+
+    assert hass.states.get("sensor.dishwasher_power").state == "1.60"
+
+    hass.states.async_set(dishwasher_mode_entity, "Cycle Active")
+    await hass.async_block_till_done()
+
+    async_fire_time_changed(hass, dt.utcnow() + timedelta(seconds=3))
+    await hass.async_block_till_done()
+
+    assert hass.states.get("sensor.dishwasher_power").state == "20.00"
+
+    async_fire_time_changed(hass, dt.utcnow() + timedelta(seconds=5))
+    await hass.async_block_till_done()
+
+    assert hass.states.get("sensor.dishwasher_power").state == "40.00"
+
+    hass.states.async_set(dishwasher_mode_entity, "Cycle Complete")
+    await hass.async_block_till_done()
+
+    assert hass.states.get("sensor.dishwasher_power").state == "9.60"
+
+    hass.states.async_set(dishwasher_mode_entity, STATE_OFF)
+    await hass.async_block_till_done()
+
+    assert hass.states.get("sensor.dishwasher_power").state == "0.00"
+
+
+async def test_calculate_standby_power(hass: HomeAssistant) -> None:
+    sensor_config = {
+        CONF_ENTITY_ID: "switch.test",
+        CONF_STANDBY_POWER: 1,
+        CONF_COMPOSITE: [
+            {
+                CONF_CONDITION: {
+                    "condition": "state",
+                    "entity_id": "switch.test",
+                    "state": STATE_OFF,
+                },
+                CONF_MULTI_SWITCH: {
+                    CONF_POWER: 5,
+                    CONF_POWER_OFF: 2,
+                    CONF_ENTITIES: [
+                        "switch.test1",
+                        "switch.test2",
+                    ],
+                },
+            },
+            {
+                CONF_FIXED: {
+                    CONF_POWER: 10,
+                },
+            },
+        ],
+    }
+
+    await run_powercalc_setup(hass, sensor_config, {})
+
+    hass.states.async_set("switch.test1", STATE_OFF)
+    hass.states.async_set("switch.test2", STATE_OFF)
+    hass.states.async_set("switch.test", STATE_OFF)
+    await hass.async_block_till_done()
+
+    assert hass.states.get("sensor.test_power").state == "4.00"
+
+    hass.states.async_set("switch.test", STATE_ON)
+    await hass.async_block_till_done()
+
+    assert hass.states.get("sensor.test_power").state == "10.00"
+
+
+async def test_calculate_standby_power2(hass: HomeAssistant) -> None:
+    sensor_config = {
+        CONF_ENTITY_ID: "switch.test",
+        CONF_STANDBY_POWER: 1,
+        CONF_COMPOSITE: [
+            {
+                CONF_CONDITION: {
+                    "condition": "state",
+                    "entity_id": "switch.test",
+                    "state": STATE_OFF,
+                },
+                CONF_FIXED: {
+                    CONF_POWER: 5,
+                },
+            },
+            {
+                CONF_MULTI_SWITCH: {
+                    CONF_POWER: 5,
+                    CONF_POWER_OFF: 2,
+                    CONF_ENTITIES: [
+                        "switch.test1",
+                        "switch.test2",
+                    ],
+                },
+            },
+        ],
+    }
+
+    await run_powercalc_setup(hass, sensor_config, {})
+
+    hass.states.async_set("switch.test", STATE_OFF)
+    await hass.async_block_till_done()
+
+    assert hass.states.get("sensor.test_power").state == "1.00"

tests/test_discovery.py

@@ -275,7 +275,7 @@ async def test_autodiscover_skips_unsupported_domains(
     mock_entity_with_model_information: MockEntityWithModel,
 ) -> None:
     mock_entity_with_model_information(
-        "media_player.test",
+        "device_tracker.test",
         "signify",
         "LCT010",
     )

tests/testing_config/powercalc/playbooks/composite/dishwasher.csv

@@ -0,0 +1,2 @@
+2,20
+4,40