Implement logic for repeat and cleanup tests

bramstroker · Aug 6, 2023 · 46dea89 · 46dea89
1 parent 526b9eb
commit 46dea89
Show file tree

Hide file tree

Showing 3 changed files with 83 additions and 92 deletions.
diff --git a/custom_components/powercalc/strategy/playbook.py b/custom_components/powercalc/strategy/playbook.py
@@ -142,7 +142,7 @@ async def _load_playbook(self, playbook_id: str) -> Playbook:
             )
 
         with open(file_path) as csv_file:
-            queue = PlaybookQueue()
+            queue = PlaybookQueue(repeat=self._repeat)
 
             csv_reader = csv.reader(csv_file)
             for row in csv_reader:
@@ -154,14 +154,18 @@ async def _load_playbook(self, playbook_id: str) -> Playbook:
 
 
 class PlaybookQueue:
-    def __init__(self) -> None:
+    def __init__(self, repeat: bool = False) -> None:
         self._entries: deque[PlaybookEntry] = deque()
+        self._repeat = repeat
 
     def enqueue(self, entry: PlaybookEntry) -> None:
         self._entries.append(entry)
 
     def dequeue(self) -> PlaybookEntry:
-        return self._entries.popleft()
+        entry = self._entries.popleft()
+        if self._repeat:
+            self.enqueue(entry)
+        return entry
 
     def __len__(self) -> int:
         return len(self._entries)

diff --git a/tests/strategy/test_playbook.py b/tests/strategy/test_playbook.py
@@ -16,6 +16,7 @@
 from custom_components.powercalc.const import (
     CONF_PLAYBOOK,
     CONF_PLAYBOOKS,
+    CONF_REPEAT,
     CONF_STANDBY_POWER,
     DOMAIN,
     DUMMY_ENTITY_ID,
@@ -30,6 +31,8 @@
     run_powercalc_setup,
 )
 
+POWER_SENSOR_ID = "sensor.test_power"
+
 
 async def test_activate_playbook_service(hass: HomeAssistant) -> None:
     hass.config.config_dir = get_test_config_dir()
@@ -46,35 +49,14 @@ async def test_activate_playbook_service(hass: HomeAssistant) -> None:
         },
     )
 
-    await hass.services.async_call(
-        DOMAIN,
-        SERVICE_ACTIVATE_PLAYBOOK,
-        {ATTR_ENTITY_ID: "sensor.test_power", "playbook_id": "playbook1"},
-        blocking=True,
-    )
-    await hass.async_block_till_done()
-
     assert hass.states.get("sensor.test_power").state == "0.00"
 
-    async_fire_time_changed(hass, dt.utcnow() + timedelta(seconds=1.5))
-    await hass.async_block_till_done()
-
-    assert hass.states.get("sensor.test_power").state == "20.50"
+    await _activate_playbook(hass, "playbook1")
 
-    async_fire_time_changed(hass, dt.utcnow() + timedelta(seconds=3))
-    await hass.async_block_till_done()
-
-    assert hass.states.get("sensor.test_power").state == "40.00"
-
-    async_fire_time_changed(hass, dt.utcnow() + timedelta(seconds=4.5))
-    await hass.async_block_till_done()
-
-    assert hass.states.get("sensor.test_power").state == "60.00"
-
-    async_fire_time_changed(hass, dt.utcnow() + timedelta(seconds=6.5))
-    await hass.async_block_till_done()
-
-    assert hass.states.get("sensor.test_power").state == "20.20"
+    await elapse_and_assert_power(hass, 1.5, "20.50")
+    await elapse_and_assert_power(hass, 3, "40.00")
+    await elapse_and_assert_power(hass, 4.5, "60.00")
+    await elapse_and_assert_power(hass, 6.5, "20.20")
 
 
 async def test_stop_playbook_service(hass: HomeAssistant) -> None:
@@ -93,52 +75,24 @@ async def test_stop_playbook_service(hass: HomeAssistant) -> None:
     )
 
     # Calling stop on a non running playbook should not raise an error.
-    await hass.services.async_call(
-        DOMAIN,
-        SERVICE_STOP_PLAYBOOK,
-        {ATTR_ENTITY_ID: "sensor.test_power"},
-        blocking=True,
-    )
-    await hass.async_block_till_done()
-
-    await hass.services.async_call(
-        DOMAIN,
-        SERVICE_ACTIVATE_PLAYBOOK,
-        {ATTR_ENTITY_ID: "sensor.test_power", "playbook_id": "playbook1"},
-        blocking=True,
-    )
-    await hass.async_block_till_done()
-
-    async_fire_time_changed(hass, dt.utcnow() + timedelta(seconds=1.5))
-    await hass.async_block_till_done()
-
-    assert hass.states.get("sensor.test_power").state == "20.50"
+    await _stop_playbook(hass)
 
-    await hass.services.async_call(
-        DOMAIN,
-        SERVICE_STOP_PLAYBOOK,
-        {ATTR_ENTITY_ID: "sensor.test_power"},
-        blocking=True,
-    )
-    await hass.async_block_till_done()
+    await _activate_playbook(hass, "playbook1")
 
-    async_fire_time_changed(hass, dt.utcnow() + timedelta(seconds=3))
-    await hass.async_block_till_done()
-
-    assert hass.states.get("sensor.test_power").state == "20.50"
+    await elapse_and_assert_power(hass, 1.5, "20.50")
 
-    async_fire_time_changed(hass, dt.utcnow() + timedelta(seconds=4.5))
-    await hass.async_block_till_done()
+    await _stop_playbook(hass)
 
-    assert hass.states.get("sensor.test_power").state == "20.50"
+    await elapse_and_assert_power(hass, 3, "20.50")
+    await elapse_and_assert_power(hass, 4.5, "20.50")
 
 
 async def test_turn_off_stops_running_playbook(hass: HomeAssistant) -> None:
     hass.config.config_dir = get_test_config_dir()
     await run_powercalc_setup(
         hass,
         {
-            CONF_ENTITY_ID: "switch.washing_machine",
+            CONF_ENTITY_ID: "switch.test",
             CONF_STANDBY_POWER: 0.5,
             CONF_PLAYBOOK: {
                 CONF_PLAYBOOKS: {
@@ -148,24 +102,15 @@ async def test_turn_off_stops_running_playbook(hass: HomeAssistant) -> None:
         },
     )
 
-    hass.states.async_set("switch.washing_machine", STATE_ON)
-    await hass.async_block_till_done()
-
-    await hass.services.async_call(
-        DOMAIN,
-        SERVICE_ACTIVATE_PLAYBOOK,
-        {ATTR_ENTITY_ID: "sensor.washing_machine_power", "playbook_id": "playbook1"},
-        blocking=True,
-    )
+    hass.states.async_set("switch.test", STATE_ON)
     await hass.async_block_till_done()
 
-    hass.states.async_set("switch.washing_machine", STATE_OFF)
-    await hass.async_block_till_done()
+    await _activate_playbook(hass, "playbook1")
 
-    async_fire_time_changed(hass, dt.utcnow() + timedelta(seconds=3))
+    hass.states.async_set("switch.test", STATE_OFF)
     await hass.async_block_till_done()
 
-    assert hass.states.get("sensor.washing_machine_power").state == "0.50"
+    await elapse_and_assert_power(hass, 3, "0.50")
 
 
 async def test_services_raises_error_on_non_playbook_sensor(
@@ -179,21 +124,34 @@ async def test_services_raises_error_on_non_playbook_sensor(
     await hass.async_block_till_done()
 
     with pytest.raises(HomeAssistantError):
-        await hass.services.async_call(
-            DOMAIN,
-            SERVICE_ACTIVATE_PLAYBOOK,
-            {ATTR_ENTITY_ID: "sensor.test_power", "playbook_id": "playbook1"},
-            blocking=True,
-        )
-        await hass.async_block_till_done()
-
-        await hass.services.async_call(
-            DOMAIN,
-            SERVICE_STOP_PLAYBOOK,
-            {ATTR_ENTITY_ID: "sensor.test_power"},
-            blocking=True,
-        )
-        await hass.async_block_till_done()
+        await _activate_playbook(hass, "playbook1")
+        await _stop_playbook(hass)
+
+
+async def test_repeat(hass: HomeAssistant) -> None:
+    hass.config.config_dir = get_test_config_dir()
+    await run_powercalc_setup(
+        hass,
+        {
+            CONF_ENTITY_ID: DUMMY_ENTITY_ID,
+            CONF_NAME: "Test",
+            CONF_PLAYBOOK: {
+                CONF_PLAYBOOKS: {
+                    "playbook": "test2.csv",
+                },
+                CONF_REPEAT: True,
+            },
+        },
+    )
+
+    await _activate_playbook(hass, "playbook")
+
+    assert hass.states.get("sensor.test_power").state == "0.00"
+
+    await elapse_and_assert_power(hass, 2, "20.00")
+    await elapse_and_assert_power(hass, 4, "40.00")
+    await elapse_and_assert_power(hass, 6, "20.00")
+    await elapse_and_assert_power(hass, 8, "40.00")
 
 
 async def test_exception_when_providing_unknown_playbook(hass: HomeAssistant) -> None:
@@ -217,3 +175,30 @@ async def test_lazy_load_playbook(hass: HomeAssistant) -> None:
     strategy = PlaybookStrategy(hass, {CONF_PLAYBOOKS: {"program1": "test.csv"}})
     await strategy.activate_playbook("program1")
     await strategy.activate_playbook("program1")
+
+
+async def elapse_and_assert_power(hass: HomeAssistant, seconds: float, expected_power: str) -> None:
+    async_fire_time_changed(hass, dt.utcnow() + timedelta(seconds=seconds))
+    await hass.async_block_till_done()
+
+    assert hass.states.get(POWER_SENSOR_ID).state == expected_power
+
+
+async def _activate_playbook(hass: HomeAssistant, playbook_id: str) -> None:
+    await hass.services.async_call(
+        DOMAIN,
+        SERVICE_ACTIVATE_PLAYBOOK,
+        {ATTR_ENTITY_ID: POWER_SENSOR_ID, "playbook_id": playbook_id},
+        blocking=True,
+    )
+    await hass.async_block_till_done()
+
+
+async def _stop_playbook(hass: HomeAssistant) -> None:
+    await hass.services.async_call(
+        DOMAIN,
+        SERVICE_STOP_PLAYBOOK,
+        {ATTR_ENTITY_ID: POWER_SENSOR_ID},
+        blocking=True,
+    )
+    await hass.async_block_till_done()
diff --git a/tests/testing_config/powercalc/playbooks/test2.csv b/tests/testing_config/powercalc/playbooks/test2.csv
@@ -0,0 +1,2 @@
+2,20
+4,40