From b48a3b288389216fb6f84ad68bcdbb2981885a45 Mon Sep 17 00:00:00 2001 From: Bram Gerritsen Date: Sat, 5 Oct 2024 10:47:11 +0200 Subject: [PATCH] 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 | 4 +- .../powercalc/strategy/composite.py | 64 ++++++- .../powercalc/strategy/multi_switch.py | 2 +- .../powercalc/strategy/playbook.py | 5 + tests/strategy/test_composite.py | 170 +++++++++++++++++- tests/test_discovery.py | 2 +- .../playbooks/composite/dishwasher.csv | 2 + 7 files changed, 238 insertions(+), 11 deletions(-) create mode 100644 tests/testing_config/powercalc/playbooks/composite/dishwasher.csv diff --git a/custom_components/powercalc/sensors/power.py b/custom_components/powercalc/sensors/power.py index 851f3c27e..2e8a3cf14 100644 --- a/custom_components/powercalc/sensors/power.py +++ b/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: diff --git a/custom_components/powercalc/strategy/composite.py b/custom_components/powercalc/strategy/composite.py index a2fe109f3..57a7bce2a 100644 --- a/custom_components/powercalc/strategy/composite.py +++ b/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) diff --git a/custom_components/powercalc/strategy/multi_switch.py b/custom_components/powercalc/strategy/multi_switch.py index 86754dbf1..9664db898 100644 --- a/custom_components/powercalc/strategy/multi_switch.py +++ b/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: diff --git a/custom_components/powercalc/strategy/playbook.py b/custom_components/powercalc/strategy/playbook.py index c930c9bf3..a9cfb7358 100644 --- a/custom_components/powercalc/strategy/playbook.py +++ b/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: diff --git a/tests/strategy/test_composite.py b/tests/strategy/test_composite.py index 5a05b9aaf..b6f6c8cd0 100644 --- a/tests/strategy/test_composite.py +++ b/tests/strategy/test_composite.py @@ -1,7 +1,11 @@ +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, @@ -9,7 +13,8 @@ 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, @@ -17,9 +22,15 @@ 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" diff --git a/tests/test_discovery.py b/tests/test_discovery.py index c5114c0de..90236c423 100644 --- a/tests/test_discovery.py +++ b/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", ) diff --git a/tests/testing_config/powercalc/playbooks/composite/dishwasher.csv b/tests/testing_config/powercalc/playbooks/composite/dishwasher.csv new file mode 100644 index 000000000..4b7f85bdc --- /dev/null +++ b/tests/testing_config/powercalc/playbooks/composite/dishwasher.csv @@ -0,0 +1,2 @@ +2,20 +4,40