Add service to turn on the water heater for an amount of time

custom_components/wundasmart/services.yaml

@@ -0,0 +1,15 @@
+hw_boost:
+  name: Boost hot water
+  description: Turns on the water heater for a specific amount of time.
+  target:
+    entity:
+      integration: wundasmart
+      domain: water_heater
+  fields:
+    duration:
+      name: Duration
+      description: Time before the water heater turns off.
+      required: true
+      advanced: false
+      example: '00:30:00'
+      default: '00:30:00'

custom_components/wundasmart/water_heater.py

@@ -2,7 +2,10 @@
 from __future__ import annotations
 
 import logging
+import math
 from typing import Any
+from aiohttp import ClientSession
+from datetime import timedelta
 
 from homeassistant.components.water_heater import (
     WaterHeaterEntity,
@@ -15,12 +18,12 @@
     CONF_PASSWORD,
     TEMP_CELSIUS,
 )
-from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers import aiohttp_client
-from homeassistant.helpers.entity import DeviceInfo
+from homeassistant.core import HomeAssistant, ServiceCall, callback
+from homeassistant.helpers import aiohttp_client, entity_platform
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
-from aiohttp import ClientSession
+import homeassistant.helpers.config_validation as cv
+import voluptuous as vol
 
 from . import WundasmartDataUpdateCoordinator
 from .pywundasmart import send_command
@@ -66,6 +69,17 @@ async def async_setup_entry(
         for wunda_id, device in coordinator.data.items() if device.get("device_type") == "wunda" and "device_name" in device
     )
 
+    platform = entity_platform.current_platform.get()
+    assert platform
+
+    platform.async_register_entity_service(
+        "hw_boost",
+        {
+            vol.Required("duration"): cv.positive_time_period
+        },
+        "async_set_boost",
+    )
+
 
 class Device(CoordinatorEntity[WundasmartDataUpdateCoordinator], WaterHeaterEntity):
     """Representation of an Wundasmart water heater."""
@@ -161,3 +175,14 @@ async def async_set_operation_mode(self, operation_mode: str) -> None:
 
         # Fetch the updated state
         await self.coordinator.async_request_refresh()
+
+    async def async_set_boost(self, duration: timedelta):
+        seconds = int((duration.days * 24 * 3600) + math.ceil(duration.seconds))
+        if seconds > 0:
+            await send_command(self._session, self._wunda_ip, self._wunda_user, self._wunda_pass, params={
+                "cmd": 3,
+                "hw_boost_time": seconds
+            })
+
+        # Fetch the updated state
+        await self.coordinator.async_request_refresh()

tests/test_water_heater.py

@@ -46,3 +46,29 @@ async def test_water_header(hass: HomeAssistant, config):
 
         assert state
         assert state.state == STATE_AUTO_OFF
+
+
+async def test_water_header_boost(hass: HomeAssistant, config):
+    entry = MockConfigEntry(domain=DOMAIN, data=config)
+    entry.add_to_hass(hass)
+
+    # Test setup of water heater entity fetches initial state
+    data = json.loads(load_fixture("test_get_devices1.json"))
+    with patch("custom_components.wundasmart.get_devices", return_value=data), \
+            patch("custom_components.wundasmart.water_heater.send_command", return_value=None) as mock:
+        await hass.config_entries.async_setup(entry.entry_id)
+        await hass.async_block_till_done()
+
+        state = hass.states.get("water_heater.smart_hubswitch")
+        assert state
+
+        await hass.services.async_call("wundasmart", "hw_boost", {
+            "entity_id": "water_heater.smart_hubswitch",
+            "duration": "00:10:00"
+        })
+        await hass.async_block_till_done()
+
+        # Check send_command was called correctly
+        assert mock.call_count == 1
+        assert mock.call_args.kwargs["params"]["cmd"] == 3
+        assert mock.call_args.kwargs["params"]["hw_boost_time"] == 600