Add hot water services and more operation modes (#12)

- Add services to turn on the water heater on and off for an amount of time
- Add a load more operation modes for standard boost/off durations

custom_components/wundasmart/services.yaml

@@ -0,0 +1,34 @@
+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'
+      selector:
+        time:
+hw_off:
+  name: Turn off hot water
+  description: Turns the water heater off for a specific amount of time.
+  target:
+    entity:
+      integration: wundasmart
+      domain: water_heater
+  fields:
+    duration:
+      name: Duration
+      description: Time to turn the water heater off for.
+      required: true
+      advanced: false
+      example: '00:30:00'
+      default: '00:30:00'
+      selector:
+        time:

custom_components/wundasmart/strings.json

@@ -31,6 +31,25 @@
           }
         }
       }
+    },
+    "water_heater": {
+      "wundasmart": {
+        "state": {
+          "auto_on": "On (Auto)",
+          "auto_off": "Off (Auto)",
+          "boost_on": "On",
+          "boost_off": "Off",
+          "auto": "Auto",
+          "boost_30": "Boost (30 mins)",
+          "boost_60": "Boost (1 hour)",
+          "boost_90": "Boost (1.5 hours)",
+          "boost_120": "Boost (2 hours)",
+          "off_30": "Off (30 mins)",
+          "off_60": "Off (1 hour)",
+          "off_90": "Off (1.5 hours)",
+          "off_120": "Off (2 hours)"
+        }
+      }
     }
   }
 }

custom_components/wundasmart/translations/en.json

@@ -31,6 +31,25 @@
                     }
                 }
             }
+        },
+        "water_heater": {
+            "wundasmart": {
+                "state": {
+                    "auto_on": "On (Auto)",
+                    "auto_off": "Off (Auto)",
+                    "boost_on": "On",
+                    "boost_off": "Off",
+                    "auto": "Auto",
+                    "boost_30": "Boost (30 mins)",
+                    "boost_60": "Boost (1 hour)",
+                    "boost_90": "Boost (1.5 hours)",
+                    "boost_120": "Boost (2 hours)",
+                    "off_30": "Off (30 mins)",
+                    "off_60": "Off (1 hour)",
+                    "off_90": "Off (1.5 hours)",
+                    "off_120": "Off (2 hours)"
+                }
+            }
         }
     }
-}
+}

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
@@ -30,18 +33,43 @@
 
 SUPPORTED_FEATURES = WaterHeaterEntityFeature.ON_OFF | WaterHeaterEntityFeature.OPERATION_MODE
 
-STATE_AUTO_ON = "On (Auto)"
-STATE_AUTO_OFF = "Off (Auto)"
-STATE_BOOST_ON = "On (Boost)"
-STATE_BOOST_OFF = "Off (Manual)"
-STATE_AUTO = "Auto"
-
-HW_BOOST_TIME = 60 * 30  # boost for 30 minutes
-HW_OFF_TIME = 60 * 60  # switch off for 1 hour
-
-OPERATION_SET_AUTO = "Auto"
-OPERATION_BOOST_ON = "Boost (30 mins)"
-OPERATION_BOOST_OFF = "Off (1 hour)"
+STATE_AUTO_ON = "auto_on"
+STATE_AUTO_OFF = "auto_off"
+STATE_BOOST_ON = "boost_on"
+STATE_BOOST_OFF = "boost_off"
+
+OPERATION_SET_AUTO = "auto"
+OPERATION_BOOST_30 = "boost_30"
+OPERATION_BOOST_60 = "boost_60"
+OPERATION_BOOST_90 = "boost_90"
+OPERATION_BOOST_120 = "boost_120"
+OPERATION_OFF_30 = "off_30"
+OPERATION_OFF_60 = "off_60"
+OPERATION_OFF_90 = "off_90"
+OPERATION_OFF_120 = "off_90"
+
+HW_BOOST_OPERATIONS = {
+    OPERATION_BOOST_30,
+    OPERATION_BOOST_60,
+    OPERATION_BOOST_90,
+    OPERATION_BOOST_120
+}
+
+HW_OFF_OPERATIONS = {
+    OPERATION_OFF_30,
+    OPERATION_OFF_60,
+    OPERATION_OFF_90,
+    OPERATION_OFF_120
+}
+
+
+def _split_operation(key):
+    """Return (operation prefix, duration in seconds)"""
+    if "_" in key:
+        key, duration = key.split("_", 1)
+        if duration.isdigit():
+            return key, int(duration) * 60
+    return key, 0
 
 
 async def async_setup_entry(
@@ -66,15 +94,33 @@ 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",
+    )
+
+    platform.async_register_entity_service(
+        "hw_off",
+        {
+            vol.Required("duration"): cv.positive_time_period
+        },
+        "async_set_off",
+    )
+
+
+
+
 
 class Device(CoordinatorEntity[WundasmartDataUpdateCoordinator], WaterHeaterEntity):
     """Representation of an Wundasmart water heater."""
 
-    _attr_operation_list = [
-        OPERATION_SET_AUTO,
-        OPERATION_BOOST_ON,
-        OPERATION_BOOST_OFF
-    ]
+    _attr_operation_list = list(sorted({ OPERATION_SET_AUTO } | HW_BOOST_OPERATIONS | HW_OFF_OPERATIONS, key=_split_operation))
     _attr_supported_features = WaterHeaterEntityFeature.OPERATION_MODE
     _attr_temperature_unit = TEMP_CELSIUS
     _attr_translation_key = DOMAIN
@@ -141,23 +187,48 @@ async def async_added_to_hass(self) -> None:
         self._handle_coordinator_update()
 
     async def async_set_operation_mode(self, operation_mode: str) -> None:
-        if operation_mode == OPERATION_BOOST_OFF:
-            await send_command(self._session, self._wunda_ip, self._wunda_user, self._wunda_pass, params={
-                "cmd": 3,
-                "hw_off_time": HW_OFF_TIME
-            })
-        elif operation_mode == OPERATION_BOOST_ON:
+        if operation_mode:
+            if operation_mode in HW_OFF_OPERATIONS:
+                _, duration = _split_operation(operation_mode)
+                await send_command(self._session, self._wunda_ip, self._wunda_user, self._wunda_pass, params={
+                    "cmd": 3,
+                    "hw_off_time": duration
+                })
+            elif operation_mode in HW_BOOST_OPERATIONS:
+                _, duration = _split_operation(operation_mode)
+                await send_command(self._session, self._wunda_ip, self._wunda_user, self._wunda_pass, params={
+                    "cmd": 3,
+                    "hw_boost_time": duration
+                })
+            elif operation_mode == OPERATION_SET_AUTO:
+                await send_command(self._session, self._wunda_ip, self._wunda_user, self._wunda_pass, params={
+                    "cmd": 3,
+                    "hw_boost_time": 0
+                })
+            else:
+                raise NotImplementedError(f"Unsupported operation mode {operation_mode}")
+
+        # 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": HW_BOOST_TIME
+                "hw_boost_time": seconds
             })
-        elif operation_mode == OPERATION_SET_AUTO:
+
+        # Fetch the updated state
+        await self.coordinator.async_request_refresh()
+
+    async def async_set_off(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": 0
+                "hw_off_time": seconds
             })
-        else:
-            raise NotImplementedError(f"Unsupported operation mode {operation_mode}")
 
         # Fetch the updated state
         await self.coordinator.async_request_refresh()

tests/test_water_heater.py

@@ -46,3 +46,64 @@ async def test_water_header(hass: HomeAssistant, config):
 
         assert state
         assert state.state == STATE_AUTO_OFF
+
+
+async def test_water_header_set_operation(hass: HomeAssistant, config):
+    entry = MockConfigEntry(domain=DOMAIN, data=config)
+    entry.add_to_hass(hass)
+
+    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("water_heater", "set_operation_mode", {
+            "entity_id": "water_heater.smart_hubswitch",
+            "operation_mode": "boost_30"
+        })
+        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"] == 1800
+
+        await hass.services.async_call("water_heater", "set_operation_mode", {
+            "entity_id": "water_heater.smart_hubswitch",
+            "operation_mode": "off_60"
+        })
+        await hass.async_block_till_done()
+
+        assert mock.call_count == 2
+        assert mock.call_args.kwargs["params"]["cmd"] == 3
+        assert mock.call_args.kwargs["params"]["hw_off_time"] == 3600
+
+
+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