Add required_features to WaterHeater entity service registrations (#141873)

homeassistant/components/atag/water_heater.py

@@ -6,6 +6,7 @@ from homeassistant.components.water_heater import (
     STATE_ECO,
     STATE_PERFORMANCE,
     WaterHeaterEntity,
+    WaterHeaterEntityFeature,
 )
 from homeassistant.const import ATTR_TEMPERATURE, STATE_OFF, Platform, UnitOfTemperature
 from homeassistant.core import HomeAssistant
@@ -32,6 +33,7 @@ class AtagWaterHeater(AtagEntity, WaterHeaterEntity):
     """Representation of an ATAG water heater."""
 
     _attr_operation_list = OPERATION_LIST
+    _attr_supported_features = WaterHeaterEntityFeature.TARGET_TEMPERATURE
     _attr_temperature_unit = UnitOfTemperature.CELSIUS
 
     @property

homeassistant/components/evohome/water_heater.py

@@ -71,6 +71,11 @@ class EvoDHW(EvoChild, WaterHeaterEntity):
     _attr_name = "DHW controller"
     _attr_icon = "mdi:thermometer-lines"
     _attr_operation_list = list(HA_STATE_TO_EVO)
+    _attr_supported_features = (
+        WaterHeaterEntityFeature.AWAY_MODE
+        | WaterHeaterEntityFeature.ON_OFF
+        | WaterHeaterEntityFeature.OPERATION_MODE
+    )
     _attr_temperature_unit = UnitOfTemperature.CELSIUS
 
     _evo_device: evo.HotWater
@@ -91,9 +96,6 @@ class EvoDHW(EvoChild, WaterHeaterEntity):
         self._attr_precision = (
             PRECISION_TENTHS if coordinator.client_v1 else PRECISION_WHOLE
         )
-        self._attr_supported_features = (
-            WaterHeaterEntityFeature.AWAY_MODE | WaterHeaterEntityFeature.OPERATION_MODE
-        )
 
     @property
     def current_operation(self) -> str | None:

homeassistant/components/hive/water_heater.py

@@ -73,7 +73,9 @@ async def async_setup_entry(
 class HiveWaterHeater(HiveEntity, WaterHeaterEntity):
     """Hive Water Heater Device."""
 
-    _attr_supported_features = WaterHeaterEntityFeature.OPERATION_MODE
+    _attr_supported_features = (
+        WaterHeaterEntityFeature.ON_OFF | WaterHeaterEntityFeature.OPERATION_MODE
+    )
     _attr_temperature_unit = UnitOfTemperature.CELSIUS
     _attr_operation_list = SUPPORT_WATER_HEATER
 

homeassistant/components/water_heater/__init__.py

@@ -112,15 +112,22 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
         SERVICE_TURN_OFF, None, "async_turn_off", [WaterHeaterEntityFeature.ON_OFF]
     )
     component.async_register_entity_service(
-        SERVICE_SET_AWAY_MODE, SET_AWAY_MODE_SCHEMA, async_service_away_mode
+        SERVICE_SET_AWAY_MODE,
+        SET_AWAY_MODE_SCHEMA,
+        async_service_away_mode,
+        [WaterHeaterEntityFeature.AWAY_MODE],
     )
     component.async_register_entity_service(
-        SERVICE_SET_TEMPERATURE, SET_TEMPERATURE_SCHEMA, async_service_temperature_set
+        SERVICE_SET_TEMPERATURE,
+        SET_TEMPERATURE_SCHEMA,
+        async_service_temperature_set,
+        [WaterHeaterEntityFeature.TARGET_TEMPERATURE],
     )
     component.async_register_entity_service(
         SERVICE_SET_OPERATION_MODE,
         SET_OPERATION_MODE_SCHEMA,
         "async_handle_set_operation_mode",
+        [WaterHeaterEntityFeature.OPERATION_MODE],
     )
 
     return True

tests/components/evohome/snapshots/test_water_heater.ambr

@@ -53,7 +53,7 @@
           'temperature': 23.0,
         }),
       }),
-      'supported_features': <WaterHeaterEntityFeature: 6>,
+      'supported_features': <WaterHeaterEntityFeature: 14>,
       'target_temp_high': None,
       'target_temp_low': None,
       'temperature': None,
@@ -100,7 +100,7 @@
           'temperature': 23.0,
         }),
       }),
-      'supported_features': <WaterHeaterEntityFeature: 6>,
+      'supported_features': <WaterHeaterEntityFeature: 14>,
       'target_temp_high': None,
       'target_temp_low': None,
       'temperature': None,

tests/components/evohome/test_water_heater.py

@@ -25,7 +25,6 @@ from homeassistant.const import (
     Platform,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import HomeAssistantError
 
 from .conftest import setup_evohome
 from .const import TEST_INSTALLS_WITH_DHW
@@ -160,8 +159,8 @@ async def test_set_away_mode(hass: HomeAssistant, evohome: EvohomeClient) -> Non
 async def test_turn_off(hass: HomeAssistant, evohome: EvohomeClient) -> None:
     """Test SERVICE_TURN_OFF of an evohome DHW zone."""
 
-    # Entity water_heater.xxx does not support this service
-    with pytest.raises(HomeAssistantError):
+    # turn_off
+    with patch("evohomeasync2.hotwater.HotWater.off") as mock_fcn:
         await hass.services.async_call(
             Platform.WATER_HEATER,
             SERVICE_TURN_OFF,
@@ -171,13 +170,15 @@ async def test_turn_off(hass: HomeAssistant, evohome: EvohomeClient) -> None:
             blocking=True,
         )
 
+        mock_fcn.assert_awaited_once_with()
+
 
 @pytest.mark.parametrize("install", TEST_INSTALLS_WITH_DHW)
 async def test_turn_on(hass: HomeAssistant, evohome: EvohomeClient) -> None:
     """Test SERVICE_TURN_ON of an evohome DHW zone."""
 
-    # Entity water_heater.xxx does not support this service
-    with pytest.raises(HomeAssistantError):
+    # turn_on
+    with patch("evohomeasync2.hotwater.HotWater.on") as mock_fcn:
         await hass.services.async_call(
             Platform.WATER_HEATER,
             SERVICE_TURN_ON,
@@ -186,3 +187,5 @@ async def test_turn_on(hass: HomeAssistant, evohome: EvohomeClient) -> None:
             },
             blocking=True,
         )
+
+        mock_fcn.assert_awaited_once_with()