Add service turn_on and turn_off service for water_heater (#94817)

homeassistant/components/demo/water_heater.py

@@ -15,6 +15,7 @@ from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
 
 SUPPORT_FLAGS_HEATER = (
     WaterHeaterEntityFeature.TARGET_TEMPERATURE
+    | WaterHeaterEntityFeature.ON_OFF
     | WaterHeaterEntityFeature.OPERATION_MODE
     | WaterHeaterEntityFeature.AWAY_MODE
 )
@@ -103,3 +104,11 @@ class DemoWaterHeater(WaterHeaterEntity):
         """Turn away mode off."""
         self._attr_is_away_mode_on = False
         self.schedule_update_ha_state()
+
+    def turn_on(self, **kwargs: Any) -> None:
+        """Turn on water heater."""
+        self.set_operation_mode("eco")
+
+    def turn_off(self, **kwargs: Any) -> None:
+        """Turn off water heater."""
+        self.set_operation_mode("off")

homeassistant/components/melcloud/water_heater.py

@@ -44,6 +44,7 @@ class AtwWaterHeater(WaterHeaterEntity):
 
     _attr_supported_features = (
         WaterHeaterEntityFeature.TARGET_TEMPERATURE
+        | WaterHeaterEntityFeature.ON_OFF
         | WaterHeaterEntityFeature.OPERATION_MODE
     )
 
@@ -72,11 +73,11 @@ class AtwWaterHeater(WaterHeaterEntity):
         """Return a device description for device registry."""
         return self._api.device_info
 
-    async def async_turn_on(self) -> None:
+    async def async_turn_on(self, **kwargs: Any) -> None:
         """Turn the entity on."""
         await self._device.set({PROPERTY_POWER: True})
 
-    async def async_turn_off(self) -> None:
+    async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn the entity off."""
         await self._device.set({PROPERTY_POWER: False})
 

homeassistant/components/water_heater/__init__.py

@@ -61,6 +61,7 @@ class WaterHeaterEntityFeature(IntFlag):
     TARGET_TEMPERATURE = 1
     OPERATION_MODE = 2
     AWAY_MODE = 4
+    ON_OFF = 8
 
 
 # These SUPPORT_* constants are deprecated as of Home Assistant 2022.5.
@@ -116,6 +117,12 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     )
     await component.async_setup(config)
 
+    component.async_register_entity_service(
+        SERVICE_TURN_ON, {}, "async_turn_on", [WaterHeaterEntityFeature.ON_OFF]
+    )
+    component.async_register_entity_service(
+        SERVICE_TURN_OFF, {}, "async_turn_off", [WaterHeaterEntityFeature.ON_OFF]
+    )
     component.async_register_entity_service(
         SERVICE_SET_AWAY_MODE, SET_AWAY_MODE_SCHEMA, async_service_away_mode
     )
@@ -294,6 +301,22 @@ class WaterHeaterEntity(Entity):
             ft.partial(self.set_temperature, **kwargs)
         )
 
+    def turn_on(self, **kwargs: Any) -> None:
+        """Turn the water heater on."""
+        raise NotImplementedError()
+
+    async def async_turn_on(self, **kwargs: Any) -> None:
+        """Turn the water heater on."""
+        await self.hass.async_add_executor_job(ft.partial(self.turn_on, **kwargs))
+
+    def turn_off(self, **kwargs: Any) -> None:
+        """Turn the water heater off."""
+        raise NotImplementedError()
+
+    async def async_turn_off(self, **kwargs: Any) -> None:
+        """Turn the water heater off."""
+        await self.hass.async_add_executor_job(ft.partial(self.turn_off, **kwargs))
+
     def set_operation_mode(self, operation_mode: str) -> None:
         """Set new target operation mode."""
         raise NotImplementedError()

homeassistant/components/water_heater/services.yaml

@@ -38,3 +38,13 @@ set_operation_mode:
       example: eco
       selector:
         text:
+
+turn_on:
+  target:
+    entity:
+      domain: water_heater
+
+turn_off:
+  target:
+    entity:
+      domain: water_heater

homeassistant/components/water_heater/strings.json

@@ -53,6 +53,14 @@
           "description": "[%key:component::water_heater::services::set_temperature::fields::operation_mode::description%]"
         }
       }
+    },
+    "turn_on": {
+      "name": "[%key:common::action::turn_on%]",
+      "description": "Turns water heater on."
+    },
+    "turn_off": {
+      "name": "[%key:common::action::turn_off%]",
+      "description": "Turns water heater off."
     }
   }
 }

tests/components/demo/test_water_heater.py

@@ -112,3 +112,19 @@ async def test_set_only_target_temp_with_convert(hass: HomeAssistant) -> None:
     await common.async_set_temperature(hass, 114, ENTITY_WATER_HEATER_CELSIUS)
     state = hass.states.get(ENTITY_WATER_HEATER_CELSIUS)
     assert state.attributes.get("temperature") == 114
+
+
+async def test_turn_on_off(hass: HomeAssistant) -> None:
+    """Test turn on and off."""
+    state = hass.states.get(ENTITY_WATER_HEATER)
+    assert state.attributes.get("temperature") == 119
+    assert state.attributes.get("away_mode") == "off"
+    assert state.attributes.get("operation_mode") == "eco"
+
+    await common.async_turn_off(hass, ENTITY_WATER_HEATER)
+    state = hass.states.get(ENTITY_WATER_HEATER)
+    assert state.attributes.get("operation_mode") == "off"
+
+    await common.async_turn_on(hass, ENTITY_WATER_HEATER)
+    state = hass.states.get(ENTITY_WATER_HEATER)
+    assert state.attributes.get("operation_mode") == "eco"

tests/components/water_heater/common.py

@@ -11,8 +11,11 @@ from homeassistant.components.water_heater import (
     SERVICE_SET_AWAY_MODE,
     SERVICE_SET_OPERATION_MODE,
     SERVICE_SET_TEMPERATURE,
+    SERVICE_TURN_OFF,
+    SERVICE_TURN_ON,
 )
 from homeassistant.const import ATTR_ENTITY_ID, ATTR_TEMPERATURE, ENTITY_MATCH_ALL
+from homeassistant.core import HomeAssistant
 
 
 async def async_set_away_mode(hass, away_mode, entity_id=ENTITY_MATCH_ALL):
@@ -54,3 +57,25 @@ async def async_set_operation_mode(hass, operation_mode, entity_id=ENTITY_MATCH_
     await hass.services.async_call(
         DOMAIN, SERVICE_SET_OPERATION_MODE, data, blocking=True
     )
+
+
+async def async_turn_on(hass: HomeAssistant, entity_id: str = ENTITY_MATCH_ALL) -> None:
+    """Turn all or specified water_heater devices on."""
+    data = {}
+
+    if entity_id:
+        data[ATTR_ENTITY_ID] = entity_id
+
+    await hass.services.async_call(DOMAIN, SERVICE_TURN_ON, data, blocking=True)
+
+
+async def async_turn_off(
+    hass: HomeAssistant, entity_id: str = ENTITY_MATCH_ALL
+) -> None:
+    """Turn all or specified water_heater devices off."""
+    data = {}
+
+    if entity_id:
+        data[ATTR_ENTITY_ID] = entity_id
+
+    await hass.services.async_call(DOMAIN, SERVICE_TURN_OFF, data, blocking=True)

tests/components/water_heater/test_init.py

@@ -0,0 +1,102 @@
+"""The tests for the water heater component."""
+from __future__ import annotations
+
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+import voluptuous as vol
+
+from homeassistant.components.water_heater import (
+    SET_TEMPERATURE_SCHEMA,
+    WaterHeaterEntity,
+    WaterHeaterEntityFeature,
+)
+from homeassistant.core import HomeAssistant
+
+from tests.common import async_mock_service
+
+
+async def test_set_temp_schema_no_req(
+    hass: HomeAssistant, caplog: pytest.LogCaptureFixture
+) -> None:
+    """Test the set temperature schema with missing required data."""
+    domain = "climate"
+    service = "test_set_temperature"
+    schema = SET_TEMPERATURE_SCHEMA
+    calls = async_mock_service(hass, domain, service, schema)
+
+    data = {"hvac_mode": "off", "entity_id": ["climate.test_id"]}
+    with pytest.raises(vol.Invalid):
+        await hass.services.async_call(domain, service, data)
+    await hass.async_block_till_done()
+
+    assert len(calls) == 0
+
+
+async def test_set_temp_schema(
+    hass: HomeAssistant, caplog: pytest.LogCaptureFixture
+) -> None:
+    """Test the set temperature schema with ok required data."""
+    domain = "water_heater"
+    service = "test_set_temperature"
+    schema = SET_TEMPERATURE_SCHEMA
+    calls = async_mock_service(hass, domain, service, schema)
+
+    data = {
+        "temperature": 20.0,
+        "operation_mode": "gas",
+        "entity_id": ["water_heater.test_id"],
+    }
+    await hass.services.async_call(domain, service, data)
+    await hass.async_block_till_done()
+
+    assert len(calls) == 1
+    assert calls[-1].data == data
+
+
+class MockWaterHeaterEntity(WaterHeaterEntity):
+    """Mock water heater device to use in tests."""
+
+    _attr_operation_list: list[str] = ["off", "heat_pump", "gas"]
+    _attr_operation = "heat_pump"
+    _attr_supported_features = WaterHeaterEntityFeature.ON_OFF
+
+
+async def test_sync_turn_on(hass: HomeAssistant) -> None:
+    """Test if async turn_on calls sync turn_on."""
+    water_heater = MockWaterHeaterEntity()
+    water_heater.hass = hass
+
+    # Test with turn_on method defined
+    setattr(water_heater, "turn_on", MagicMock())
+    await water_heater.async_turn_on()
+
+    # pylint: disable-next=no-member
+    assert water_heater.turn_on.call_count == 1
+
+    # Test with async_turn_on method defined
+    setattr(water_heater, "async_turn_on", AsyncMock())
+    await water_heater.async_turn_on()
+
+    # pylint: disable-next=no-member
+    assert water_heater.async_turn_on.call_count == 1
+
+
+async def test_sync_turn_off(hass: HomeAssistant) -> None:
+    """Test if async turn_off calls sync turn_off."""
+    water_heater = MockWaterHeaterEntity()
+    water_heater.hass = hass
+
+    # Test with turn_off method defined
+    setattr(water_heater, "turn_off", MagicMock())
+    await water_heater.async_turn_off()
+
+    # pylint: disable-next=no-member
+    assert water_heater.turn_off.call_count == 1
+
+    # Test with async_turn_off method defined
+    setattr(water_heater, "async_turn_off", AsyncMock())
+    await water_heater.async_turn_off()
+
+    # pylint: disable-next=no-member
+    assert water_heater.async_turn_off.call_count == 1