From 26678277a8deb09158fc97a86dfac22a07586a99 Mon Sep 17 00:00:00 2001
From: Rob Ferrer <rob@ferrer.co.uk>
Date: Sat, 4 Nov 2023 15:09:17 +0000
Subject: [PATCH] Add stop boost service

Closes #172
---
 README.md                                |  1 +
 custom_components/myenergi/entity.py     |  6 ++++++
 custom_components/myenergi/select.py     |  5 +++++
 custom_components/myenergi/services.yaml |  8 ++++++++
 tests/conftest.py                        |  7 +++++++
 tests/test_services.py                   | 20 ++++++++++++++++++++
 6 files changed, 47 insertions(+)

diff --git a/README.md b/README.md
index fdcee1a..037a410 100644
--- a/README.md
+++ b/README.md
@@ -45,6 +45,7 @@ It will create HA devices depending on what you have installed:
   - Minumum green level number input; how much power must be sourced from green sources (local generation) to do diversion charging
   - Service to start boost (provide boost amount in kWh as parameter)
   - Service to start smart boost (provide boost amount in kWh and desired finished time as paramters)
+  - Service to stop boost
 
 - Eddi
 
diff --git a/custom_components/myenergi/entity.py b/custom_components/myenergi/entity.py
index 97b7fb1..c85e9e8 100644
--- a/custom_components/myenergi/entity.py
+++ b/custom_components/myenergi/entity.py
@@ -64,6 +64,12 @@ async def start_smart_boost(self, amount: float, when: str) -> None:
         await self.device.start_smart_boost(amount, when)
         self.schedule_update_ha_state()
 
+    async def stop_boost(self) -> None:
+        _LOGGER.debug("Stop boost called")
+        """Stop boost"""
+        await self.device.stop_boost()
+        self.schedule_update_ha_state()
+
 
 class MyenergiHub(CoordinatorEntity):
     def __init__(self, coordinator, config_entry, meta):
diff --git a/custom_components/myenergi/select.py b/custom_components/myenergi/select.py
index 47598f0..c457e56 100644
--- a/custom_components/myenergi/select.py
+++ b/custom_components/myenergi/select.py
@@ -56,6 +56,11 @@ async def async_setup_entry(hass, entry, async_add_devices):
                 SMART_BOOST_SCHEMA,
                 "start_smart_boost",
             )
+            platform.async_register_entity_service(
+                "myenergi_stop_boost",
+                {},
+                "stop_boost",
+            )
             devices.append(ZappiChargeModeSelect(coordinator, device, entry))
         elif device.kind == "eddi":
             platform.async_register_entity_service(
diff --git a/custom_components/myenergi/services.yaml b/custom_components/myenergi/services.yaml
index c494e9c..cbe63d4 100644
--- a/custom_components/myenergi/services.yaml
+++ b/custom_components/myenergi/services.yaml
@@ -63,3 +63,11 @@ myenergi_smart_boost:
       required: true
       selector:
         time:
+myenergi_stop_boost:
+  name: Stop boost
+  description: Stop boost
+  target:
+    device:
+      model: Zappi
+    entity:
+      domain: select
diff --git a/tests/conftest.py b/tests/conftest.py
index 3561145..20f4d5f 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -138,3 +138,10 @@ def mock_eddi_device():
     """Return a mocked eddi object."""
     with patch("pymyenergi.client.Eddi.set_priority") as device:
         yield device
+
+
+@pytest.fixture
+def mock_zappi_stop_boost():
+    """Return a mocked Zappi object."""
+    with patch("pymyenergi.client.Zappi.stop_boost") as stop_boost:
+        yield stop_boost
diff --git a/tests/test_services.py b/tests/test_services.py
index d7d4925..d9b64c1 100644
--- a/tests/test_services.py
+++ b/tests/test_services.py
@@ -70,3 +70,23 @@ async def test_eddi_boost(
     await hass.async_block_till_done()
     assert mock_eddi_manual_boost.call_count == 1
     mock_eddi_manual_boost.assert_called_with("Heater 1", 44.0)
+
+
+async def test_stop_boost(
+    hass: HomeAssistant, mock_zappi_stop_boost: MagicMock
+) -> None:
+    """Verify device information includes expected details."""
+
+    await setup_mock_myenergi_config_entry(hass)
+
+    await hass.services.async_call(
+        "myenergi",
+        "myenergi_stop_boost",
+        {
+            ATTR_ENTITY_ID: TEST_ZAPPI_SELECT_CHARGE_MODE,
+        },
+        blocking=False,
+    )
+    assert mock_zappi_stop_boost.call_count == 0
+    await hass.async_block_till_done()
+    assert mock_zappi_stop_boost.call_count == 1