From 4e399680d8cc0dc10b8d346f0095e355fe8cd6c9 Mon Sep 17 00:00:00 2001
From: Bram Gerritsen <bgerritsen@emico.nl>
Date: Sun, 13 Oct 2024 17:56:54 +0200
Subject: [PATCH] fix: issue with playbook returning old values when stopped
 (#2583)

---
 .../powercalc/strategy/playbook.py            | 10 ++--
 tests/strategy/test_playbook.py               | 37 +++++++++++++++
 tests/test_sensor.py                          | 46 +++++++++++++++++++
 3 files changed, 90 insertions(+), 3 deletions(-)

diff --git a/custom_components/powercalc/strategy/playbook.py b/custom_components/powercalc/strategy/playbook.py
index a9cfb7358..7013c49e0 100644
--- a/custom_components/powercalc/strategy/playbook.py
+++ b/custom_components/powercalc/strategy/playbook.py
@@ -82,9 +82,12 @@ def set_update_callback(self, update_callback: Callable[[Decimal], None]) -> Non
         self._update_callback = update_callback
 
     async def calculate(self, entity_state: State) -> Decimal | None:
-        if self._states_trigger and entity_state.state in self._states_trigger:
-            playbook_id = self._states_trigger[entity_state.state]
-            await self.activate_playbook(playbook_id)
+        if self._states_trigger:
+            if entity_state.state in self._states_trigger:
+                playbook_id = self._states_trigger[entity_state.state]
+                await self.activate_playbook(playbook_id)
+            else:
+                await self.stop_playbook()
 
         return self._power
 
@@ -112,6 +115,7 @@ async def stop_playbook(self) -> None:
 
         _LOGGER.debug("Stopping playbook")
         self._active_playbook = None
+        self._power = Decimal(0)
         if self._cancel_timer is not None:
             self._cancel_timer()
             self._cancel_timer = None
diff --git a/tests/strategy/test_playbook.py b/tests/strategy/test_playbook.py
index 5bff6c881..c97955cee 100644
--- a/tests/strategy/test_playbook.py
+++ b/tests/strategy/test_playbook.py
@@ -318,6 +318,43 @@ async def test_multiply_factor(hass: HomeAssistant) -> None:
     await elapse_and_assert_power(hass, 2, "60.00")
 
 
+async def test_source_entity_trigger(hass: HomeAssistant) -> None:
+    hass.config.config_dir = get_test_config_dir()
+    await run_powercalc_setup(
+        hass,
+        {
+            CONF_ENTITY_ID: "switch.test",
+            CONF_NAME: "Test",
+            CONF_PLAYBOOK: {
+                CONF_PLAYBOOKS: {
+                    "on": "test2.csv",
+                },
+                CONF_STATE_TRIGGER: {
+                    STATE_ON: "on",
+                },
+            },
+        },
+    )
+
+    hass.states.async_set("switch.test", STATE_ON)
+    await hass.async_block_till_done()
+
+    assert hass.states.get(POWER_SENSOR_ID).state == "0.00"
+    await elapse_and_assert_power(hass, 2, "20.00")
+
+    hass.states.async_set("switch.test", STATE_OFF)
+    await hass.async_block_till_done()
+
+    async_fire_time_changed(hass, dt.utcnow() + timedelta(seconds=60))
+    await hass.async_block_till_done()
+
+    hass.states.async_set("switch.test", STATE_ON)
+    await hass.async_block_till_done()
+
+    assert hass.states.get(POWER_SENSOR_ID).state == "0.00"
+    await elapse_and_assert_power(hass, 2, "20.00")
+
+
 async def test_state_trigger(hass: HomeAssistant) -> None:
     hass.config.config_dir = get_test_config_dir()
     await run_powercalc_setup(
diff --git a/tests/test_sensor.py b/tests/test_sensor.py
index 0b3879438..a5852e655 100644
--- a/tests/test_sensor.py
+++ b/tests/test_sensor.py
@@ -51,6 +51,7 @@
     CONF_CREATE_GROUP,
     CONF_CREATE_UTILITY_METERS,
     CONF_ENABLE_AUTODISCOVERY,
+    CONF_ENERGY_INTEGRATION_METHOD,
     CONF_ENERGY_SENSOR_FRIENDLY_NAMING,
     CONF_ENERGY_SENSOR_NAMING,
     CONF_FIXED,
@@ -62,6 +63,7 @@
     CONF_SENSOR_TYPE,
     CONF_UTILITY_METER_TYPES,
     DOMAIN,
+    ENERGY_INTEGRATION_METHOD_LEFT,
     CalculationStrategy,
     SensorType,
 )
@@ -852,3 +854,47 @@ async def test_change_config_entry_entity_id(hass: HomeAssistant) -> None:
     power_state = hass.states.get("sensor.testentry_power")
     assert power_state.attributes.get(ATTR_SOURCE_ENTITY) == new_light_id
     assert power_state.state == "100.00"
+
+
+async def test_regression(hass: HomeAssistant) -> None:
+    #
+    # {
+    #     "entity_id": "light.standing_lamp",
+    #     "mode": "fixed",
+    #     "create_energy_sensor": true,
+    #     "create_utility_meters": false,
+    #     "fixed": {
+    #         "power": 11.0
+    #     },
+    #     "energy_integration_method": "left",
+    #     "unique_id": "pc_91b5b2b04e889549cb19effa3dea97e1",
+    #     "sensor_type": "virtual_power",
+    #     "name": "Standing lamp",
+    #     "_power_entity": "sensor.standing_lamp_power",
+    #     "_energy_entity": "sensor.standing_lamp_power"
+    # }
+
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        data={
+            CONF_SENSOR_TYPE: SensorType.VIRTUAL_POWER,
+            CONF_CREATE_ENERGY_SENSOR: True,
+            CONF_CREATE_UTILITY_METERS: False,
+            CONF_MODE: CalculationStrategy.FIXED,
+            CONF_NAME: "Standing lamp",
+            CONF_ENTITY_ID: "light.standing_lamp",
+            CONF_ENERGY_INTEGRATION_METHOD: ENERGY_INTEGRATION_METHOD_LEFT,
+            CONF_FIXED: {CONF_POWER: 11},
+            CONF_UNIQUE_ID: "pc_91b5b2b04e889549cb19effa3dea97e1",
+        },
+        unique_id="pc_91b5b2b04e889549cb19effa3dea97e1",
+    )
+    entry.add_to_hass(hass)
+
+    await run_powercalc_setup(
+        hass,
+        {},
+    )
+
+    assert hass.states.get("sensor.standing_lamp_power")
+    assert hass.states.get("sensor.standing_lamp_energy")