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

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