From 99ab2566c25a3ba40b205302065740907675bd84 Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Fri, 16 Aug 2024 12:58:05 +0200
Subject: [PATCH] Correct water heater service schemas (#124038)

* Correct water heater service schemas

* Update tests
---
 .../components/water_heater/__init__.py       | 42 ++++++-------------
 tests/components/water_heater/test_init.py    |  5 ++-
 2 files changed, 15 insertions(+), 32 deletions(-)

diff --git a/homeassistant/components/water_heater/__init__.py b/homeassistant/components/water_heater/__init__.py
index 0b54a4c1aa4..2e749735b0c 100644
--- a/homeassistant/components/water_heater/__init__.py
+++ b/homeassistant/components/water_heater/__init__.py
@@ -35,7 +35,7 @@ from homeassistant.helpers.deprecation import (
 from homeassistant.helpers.entity import Entity, EntityDescription
 from homeassistant.helpers.entity_component import EntityComponent
 from homeassistant.helpers.temperature import display_temp as show_temp
-from homeassistant.helpers.typing import ConfigType
+from homeassistant.helpers.typing import ConfigType, VolDictType
 from homeassistant.util.unit_conversion import TemperatureConverter
 
 from .const import DOMAIN
@@ -94,29 +94,17 @@ CONVERTIBLE_ATTRIBUTE = [ATTR_TEMPERATURE]
 
 _LOGGER = logging.getLogger(__name__)
 
-ON_OFF_SERVICE_SCHEMA = vol.Schema({vol.Optional(ATTR_ENTITY_ID): cv.comp_entity_ids})
-
-SET_AWAY_MODE_SCHEMA = vol.Schema(
-    {
-        vol.Optional(ATTR_ENTITY_ID): cv.comp_entity_ids,
-        vol.Required(ATTR_AWAY_MODE): cv.boolean,
-    }
-)
-SET_TEMPERATURE_SCHEMA = vol.Schema(
-    vol.All(
-        {
-            vol.Required(ATTR_TEMPERATURE, "temperature"): vol.Coerce(float),
-            vol.Optional(ATTR_ENTITY_ID): cv.comp_entity_ids,
-            vol.Optional(ATTR_OPERATION_MODE): cv.string,
-        }
-    )
-)
-SET_OPERATION_MODE_SCHEMA = vol.Schema(
-    {
-        vol.Optional(ATTR_ENTITY_ID): cv.comp_entity_ids,
-        vol.Required(ATTR_OPERATION_MODE): cv.string,
-    }
-)
+SET_AWAY_MODE_SCHEMA: VolDictType = {
+    vol.Required(ATTR_AWAY_MODE): cv.boolean,
+}
+SET_TEMPERATURE_SCHEMA: VolDictType = {
+    vol.Required(ATTR_TEMPERATURE, "temperature"): vol.Coerce(float),
+    vol.Optional(ATTR_ENTITY_ID): cv.comp_entity_ids,
+    vol.Optional(ATTR_OPERATION_MODE): cv.string,
+}
+SET_OPERATION_MODE_SCHEMA: VolDictType = {
+    vol.Required(ATTR_OPERATION_MODE): cv.string,
+}
 
 # mypy: disallow-any-generics
 
@@ -145,12 +133,6 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
         SET_OPERATION_MODE_SCHEMA,
         "async_handle_set_operation_mode",
     )
-    component.async_register_entity_service(
-        SERVICE_TURN_OFF, ON_OFF_SERVICE_SCHEMA, "async_turn_off"
-    )
-    component.async_register_entity_service(
-        SERVICE_TURN_ON, ON_OFF_SERVICE_SCHEMA, "async_turn_on"
-    )
 
     return True
 
diff --git a/tests/components/water_heater/test_init.py b/tests/components/water_heater/test_init.py
index f883cf47b19..4e0f860366c 100644
--- a/tests/components/water_heater/test_init.py
+++ b/tests/components/water_heater/test_init.py
@@ -22,6 +22,7 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import UnitOfTemperature
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ServiceValidationError
+from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from tests.common import (
@@ -42,7 +43,7 @@ async def test_set_temp_schema_no_req(
     """Test the set temperature schema with missing required data."""
     domain = "climate"
     service = "test_set_temperature"
-    schema = SET_TEMPERATURE_SCHEMA
+    schema = cv.make_entity_service_schema(SET_TEMPERATURE_SCHEMA)
     calls = async_mock_service(hass, domain, service, schema)
 
     data = {"hvac_mode": "off", "entity_id": ["climate.test_id"]}
@@ -59,7 +60,7 @@ async def test_set_temp_schema(
     """Test the set temperature schema with ok required data."""
     domain = "water_heater"
     service = "test_set_temperature"
-    schema = SET_TEMPERATURE_SCHEMA
+    schema = cv.make_entity_service_schema(SET_TEMPERATURE_SCHEMA)
     calls = async_mock_service(hass, domain, service, schema)
 
     data = {