Add support for water_heater entities

aioesphomeapi/api.proto

@@ -63,6 +63,8 @@ service APIConnection {
   rpc subscribe_voice_assistant(SubscribeVoiceAssistantRequest) returns (void) {}
 
   rpc alarm_control_panel_command (AlarmControlPanelCommandRequest) returns (void) {}
+
+  rpc water_heater_command (WaterHeaterCommandRequest) returns (void) {}
 }
 
 
@@ -1904,3 +1906,64 @@ message UpdateCommandRequest {
   fixed32 key = 1;
   UpdateCommand command = 2;
 }
+
+// ==================== WATER HEATER ====================
+enum WaterHeaterMode {
+  WATER_HEATER_MODE_OFF = 0;
+  WATER_HEATER_MODE_ECO = 1;
+  WATER_HEATER_MODE_ELECTRIC = 2;
+  WATER_HEATER_MODE_PERFORMANCE = 3;
+  WATER_HEATER_MODE_HIGH_DEMAND = 4;
+  WATER_HEATER_MODE_HEAT_PUMP = 5;
+  WATER_HEATER_MODE_GAS = 6;
+}
+message ListEntitiesWaterHeaterResponse {
+  option (id) = 119;
+  option (source) = SOURCE_SERVER;
+  option (ifdef) = "USE_WATER_HEATER";
+
+  string object_id = 1;
+  fixed32 key = 2;
+  string name = 3;
+  string unique_id = 4;
+
+  bool supports_current_temperature = 5;
+  bool supports_two_point_target_temperature = 6;
+  repeated WaterHeaterMode supported_modes = 7;
+  float visual_min_temperature = 8;
+  float visual_max_temperature = 9;
+  float visual_target_temperature_step = 10;
+  bool disabled_by_default = 11;
+  string icon = 12;
+  EntityCategory entity_category = 13;
+  float visual_current_temperature_step = 14;
+}
+message WaterHeaterStateResponse {
+  option (id) = 120;
+  option (source) = SOURCE_SERVER;
+  option (ifdef) = "USE_WATER_HEATER";
+  option (no_delay) = true;
+
+  fixed32 key = 1;
+  WaterHeaterMode mode = 2;
+  float current_temperature = 3;
+  float target_temperature = 4;
+  float target_temperature_low = 5;
+  float target_temperature_high = 6;
+}
+message WaterHeaterCommandRequest {
+  option (id) = 121;
+  option (source) = SOURCE_CLIENT;
+  option (ifdef) = "USE_WATER_HEATER";
+  option (no_delay) = true;
+
+  fixed32 key = 1;
+  bool has_mode = 2;
+  WaterHeaterMode mode = 3;
+  bool has_target_temperature = 4;
+  float target_temperature = 5;
+  bool has_target_temperature_low = 6;
+  float target_temperature_low = 7;
+  bool has_target_temperature_high = 8;
+  float target_temperature_high = 9;
+}

aioesphomeapi/client.py

@@ -76,6 +76,7 @@
     VoiceAssistantRequest,
     VoiceAssistantResponse,
     VoiceAssistantTimerEventResponse,
+    WaterHeaterCommandRequest,
 )
 from .client_callbacks import (
     on_bluetooth_connections_free_response,
@@ -133,6 +134,7 @@
     VoiceAssistantEventType,
     VoiceAssistantSubscriptionFlag,
     VoiceAssistantTimerEventType,
+    WaterHeaterMode,
     message_types_to_names,
 )
 from .model_conversions import (
@@ -1421,3 +1423,27 @@ def alarm_control_panel_command(
         if code is not None:
             req.code = code
         self._get_connection().send_message(req)
+
+    def water_heater_command(
+        self,
+        key: int,
+        mode: WaterHeaterMode | None = None,
+        target_temperature: float | None = None,
+        target_temperature_low: float | None = None,
+        target_temperature_high: float | None = None,
+    ) -> None:
+        connection = self._get_connection()
+        req = WaterHeaterCommandRequest(key=key)
+        if mode is not None:
+            req.has_mode = True
+            req.mode = mode
+        if target_temperature is not None:
+            req.has_target_temperature = True
+            req.target_temperature = target_temperature
+        if target_temperature_low is not None:
+            req.has_target_temperature_low = True
+            req.target_temperature_low = target_temperature_low
+        if target_temperature_high is not None:
+            req.has_target_temperature_high = True
+            req.target_temperature_high = target_temperature_high
+        connection.send_message(req)

aioesphomeapi/core.py

@@ -84,6 +84,7 @@
     ListEntitiesTimeResponse,
     ListEntitiesUpdateResponse,
     ListEntitiesValveResponse,
+    ListEntitiesWaterHeaterResponse,
     LockCommandRequest,
     LockStateResponse,
     MediaPlayerCommandRequest,
@@ -123,6 +124,8 @@
     VoiceAssistantRequest,
     VoiceAssistantResponse,
     VoiceAssistantTimerEventResponse,
+    WaterHeaterCommandRequest,
+    WaterHeaterStateResponse,
 )
 
 TWO_CHAR = re.compile(r".{2}")
@@ -392,4 +395,7 @@ def __init__(self, error: BluetoothGATTError) -> None:
     116: ListEntitiesUpdateResponse,
     117: UpdateStateResponse,
     118: UpdateCommandRequest,
+    119: ListEntitiesWaterHeaterResponse,
+    120: WaterHeaterStateResponse,
+    121: WaterHeaterCommandRequest,
 }

aioesphomeapi/model.py

@@ -1297,6 +1297,57 @@ class VoiceAssistantTimerEventType(APIIntEnum):
     VOICE_ASSISTANT_TIMER_FINISHED = 3
 
 
+# ==================== WATER HEATER ====================
+class WaterHeaterMode(APIIntEnum):
+    OFF = 0
+    ECO = 1
+    ELECTRIC = 2
+    PERFORMANCE = 3
+    HIGH_DEMAND = 4
+    HEAT_PUMP = 5
+    GAS = 6
+
+
+@_frozen_dataclass_decorator
+class WaterHeaterInfo(EntityInfo):
+    supports_current_temperature: bool = False
+    supports_two_point_target_temperature: bool = False
+    supported_modes: list[WaterHeaterMode] = converter_field(
+        default_factory=list, converter=WaterHeaterMode.convert_list
+    )
+    visual_min_temperature: float = converter_field(
+        default=0.0, converter=fix_float_single_double_conversion
+    )
+    visual_max_temperature: float = converter_field(
+        default=0.0, converter=fix_float_single_double_conversion
+    )
+    visual_target_temperature_step: float = converter_field(
+        default=0.0, converter=fix_float_single_double_conversion
+    )
+    visual_current_temperature_step: float = converter_field(
+        default=0.0, converter=fix_float_single_double_conversion
+    )
+
+
+@_frozen_dataclass_decorator
+class WaterHeaterState(EntityState):
+    mode: WaterHeaterMode | None = converter_field(
+        default=WaterHeaterMode.OFF, converter=WaterHeaterMode.convert
+    )
+    current_temperature: float = converter_field(
+        default=0.0, converter=fix_float_single_double_conversion
+    )
+    target_temperature: float = converter_field(
+        default=0.0, converter=fix_float_single_double_conversion
+    )
+    target_temperature_low: float = converter_field(
+        default=0.0, converter=fix_float_single_double_conversion
+    )
+    target_temperature_high: float = converter_field(
+        default=0.0, converter=fix_float_single_double_conversion
+    )
+
+
 _TYPE_TO_NAME = {
     BinarySensorInfo: "binary_sensor",
     ButtonInfo: "button",
@@ -1321,6 +1372,7 @@ class VoiceAssistantTimerEventType(APIIntEnum):
     ValveInfo: "valve",
     EventInfo: "event",
     UpdateInfo: "update",
+    WaterHeaterInfo: "water_heater",
 }
 
 

aioesphomeapi/model_conversions.py

@@ -36,6 +36,7 @@
     ListEntitiesTimeResponse,
     ListEntitiesUpdateResponse,
     ListEntitiesValveResponse,
+    ListEntitiesWaterHeaterResponse,
     LockStateResponse,
     MediaPlayerStateResponse,
     NumberStateResponse,
@@ -48,6 +49,7 @@
     TimeStateResponse,
     UpdateStateResponse,
     ValveStateResponse,
+    WaterHeaterStateResponse,
 )
 from .model import (
     AlarmControlPanelEntityState,
@@ -96,6 +98,8 @@
     UpdateState,
     ValveInfo,
     ValveState,
+    WaterHeaterInfo,
+    WaterHeaterState,
 )
 
 SUBSCRIBE_STATES_RESPONSE_TYPES: dict[Any, type[EntityState]] = {
@@ -120,6 +124,7 @@
     TimeStateResponse: TimeState,
     UpdateStateResponse: UpdateState,
     ValveStateResponse: ValveState,
+    WaterHeaterStateResponse: WaterHeaterState,
 }
 
 LIST_ENTITIES_SERVICES_RESPONSE_TYPES: dict[Any, type[EntityInfo] | None] = {
@@ -147,4 +152,5 @@
     ListEntitiesTimeResponse: TimeInfo,
     ListEntitiesUpdateResponse: UpdateInfo,
     ListEntitiesValveResponse: ValveInfo,
+    ListEntitiesWaterHeaterResponse: WaterHeaterInfo,
 }

tests/test_client.py

@@ -2,15 +2,15 @@
 
 import asyncio
 import contextlib
+from functools import partial
 import itertools
 import logging
 import socket
-from functools import partial
 from typing import Any
 from unittest.mock import AsyncMock, MagicMock, call, create_autospec, patch
 
-import pytest
 from google.protobuf import message
+import pytest
 
 from aioesphomeapi._frame_helper.plain_text import APIPlaintextFrameHelper
 from aioesphomeapi.api_pb2 import (
@@ -74,6 +74,7 @@
     VoiceAssistantRequest,
     VoiceAssistantResponse,
     VoiceAssistantTimerEventResponse,
+    WaterHeaterCommandRequest,
 )
 from aioesphomeapi.client import APIClient, BluetoothConnectionDroppedError
 from aioesphomeapi.connection import APIConnection
@@ -89,6 +90,7 @@
     BinarySensorInfo,
     BinarySensorState,
     BluetoothDeviceRequestType,
+    BluetoothGATTService as BluetoothGATTServiceModel,
     BluetoothLEAdvertisement,
     BluetoothProxyFeature,
     CameraState,
@@ -108,14 +110,10 @@
     UserService,
     UserServiceArg,
     UserServiceArgType,
-)
-from aioesphomeapi.model import BluetoothGATTService as BluetoothGATTServiceModel
-from aioesphomeapi.model import (
     VoiceAssistantAudioSettings as VoiceAssistantAudioSettingsModel,
-)
-from aioesphomeapi.model import VoiceAssistantEventType as VoiceAssistantEventModelType
-from aioesphomeapi.model import (
+    VoiceAssistantEventType as VoiceAssistantEventModelType,
     VoiceAssistantTimerEventType as VoiceAssistantTimerEventModelType,
+    WaterHeaterMode,
 )
 from aioesphomeapi.reconnect_logic import ReconnectLogic, ReconnectLogicState
 
@@ -2604,3 +2602,34 @@ async def test_calls_after_connection_closed(
 
     with pytest.raises(APIConnectionError):
         await client.update_command(1, True)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "cmd, req",
+    [
+        (
+            dict(key=1, mode=WaterHeaterMode.HEAT_PUMP),
+            dict(key=1, has_mode=True, mode=WaterHeaterMode.HEAT_PUMP),
+        ),
+        (
+            dict(key=1, target_temperature=21.0),
+            dict(key=1, has_target_temperature=True, target_temperature=21.0),
+        ),
+        (
+            dict(key=1, target_temperature_low=21.0),
+            dict(key=1, has_target_temperature_low=True, target_temperature_low=21.0),
+        ),
+        (
+            dict(key=1, target_temperature_high=21.0),
+            dict(key=1, has_target_temperature_high=True, target_temperature_high=21.0),
+        ),
+    ],
+)
+async def test_water_heater_command(
+    auth_client: APIClient, cmd: dict[str, Any], req: dict[str, Any]
+) -> None:
+    send = patch_send(auth_client)
+
+    auth_client.water_heater_command(**cmd)
+    send.assert_called_once_with(WaterHeaterCommandRequest(**req))

tests/test_model.py

@@ -43,6 +43,7 @@
     ListEntitiesTimeResponse,
     ListEntitiesUpdateResponse,
     ListEntitiesValveResponse,
+    ListEntitiesWaterHeaterResponse,
     LockStateResponse,
     MediaPlayerStateResponse,
     NumberStateResponse,
@@ -55,6 +56,7 @@
     TimeStateResponse,
     UpdateStateResponse,
     ValveStateResponse,
+    WaterHeaterStateResponse,
 )
 from aioesphomeapi.model import (
     _TYPE_TO_NAME,
@@ -122,6 +124,8 @@
     ValveInfo,
     ValveState,
     VoiceAssistantFeature,
+    WaterHeaterInfo,
+    WaterHeaterState,
     build_unique_id,
     converter_field,
 )
@@ -292,6 +296,8 @@ def test_api_version_ord():
         (Event, EventResponse),
         (UpdateInfo, ListEntitiesUpdateResponse),
         (UpdateState, UpdateStateResponse),
+        (WaterHeaterInfo, ListEntitiesWaterHeaterResponse),
+        (WaterHeaterState, WaterHeaterStateResponse),
     ],
 )
 def test_basic_pb_conversions(model, pb):
@@ -409,6 +415,7 @@ def test_user_service_conversion():
         AlarmControlPanelInfo,
         TextInfo,
         TimeInfo,
+        WaterHeaterInfo,
     ],
 )
 def test_build_unique_id(model):