Sort the water heater operation modes

custom_components/wundasmart/water_heater.py

@@ -62,19 +62,14 @@
     OPERATION_OFF_120
 }
 
-HW_BOOST_TIME = {
-    OPERATION_BOOST_30: 30 * 60,
-    OPERATION_BOOST_60: 60 * 60,
-    OPERATION_BOOST_90: 90 * 60,
-    OPERATION_BOOST_120: 120 * 60
-}
 
-HW_OFF_TIME = {
-    OPERATION_OFF_30: 30 * 60,
-    OPERATION_OFF_60: 60 * 60,
-    OPERATION_OFF_90: 90 * 60,
-    OPERATION_OFF_120: 120 * 60
-}
+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(
@@ -119,10 +114,13 @@ async def async_setup_entry(
     )
 
 
+
+
+
 class Device(CoordinatorEntity[WundasmartDataUpdateCoordinator], WaterHeaterEntity):
     """Representation of an Wundasmart water heater."""
 
-    _attr_operation_list = list({ OPERATION_SET_AUTO } | HW_BOOST_OPERATIONS | HW_OFF_OPERATIONS)
+    _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
@@ -191,14 +189,16 @@ async def async_added_to_hass(self) -> None:
     async def async_set_operation_mode(self, operation_mode: str) -> None:
         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": HW_OFF_TIME[operation_mode]
+                    "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": HW_BOOST_TIME[operation_mode]
+                    "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={

tests/test_water_heater.py

@@ -48,6 +48,41 @@ async def test_water_header(hass: HomeAssistant, config):
         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)