Add validation to water_heater set_operation mode at entity component (#111168)

* Add validation to water_heater set_operation mode at entity component

* Add final decorator

homeassistant/components/water_heater/__init__.py

@@ -23,6 +23,7 @@ from homeassistant.const import (
     UnitOfTemperature,
 )
 from homeassistant.core import HomeAssistant, ServiceCall
+from homeassistant.exceptions import ServiceValidationError
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.config_validation import (  # noqa: F401
     PLATFORM_SCHEMA,
@@ -149,7 +150,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     component.async_register_entity_service(
         SERVICE_SET_OPERATION_MODE,
         SET_OPERATION_MODE_SCHEMA,
-        "async_set_operation_mode",
+        "async_handle_set_operation_mode",
     )
     component.async_register_entity_service(
         SERVICE_TURN_OFF, ON_OFF_SERVICE_SCHEMA, "async_turn_off"
@@ -359,6 +360,36 @@ class WaterHeaterEntity(Entity, cached_properties=CACHED_PROPERTIES_WITH_ATTR_):
         """Set new target operation mode."""
         await self.hass.async_add_executor_job(self.set_operation_mode, operation_mode)
 
+    @final
+    async def async_handle_set_operation_mode(self, operation_mode: str) -> None:
+        """Handle a set target operation mode service call."""
+        if self.operation_list is None:
+            raise ServiceValidationError(
+                f"Operation mode {operation_mode} not valid for "
+                f"entity {self.entity_id}. The operation list is not defined",
+                translation_domain=DOMAIN,
+                translation_key="operation_list_not_defined",
+                translation_placeholders={
+                    "entity_id": self.entity_id,
+                    "operation_mode": operation_mode,
+                },
+            )
+        if operation_mode not in self.operation_list:
+            operation_list = ", ".join(self.operation_list)
+            raise ServiceValidationError(
+                f"Operation mode {operation_mode} not valid for "
+                f"entity {self.entity_id}. Valid "
+                f"operation modes are: {operation_list}",
+                translation_domain=DOMAIN,
+                translation_key="not_valid_operation_mode",
+                translation_placeholders={
+                    "entity_id": self.entity_id,
+                    "operation_mode": operation_mode,
+                    "operation_list": operation_list,
+                },
+            )
+        await self.async_set_operation_mode(operation_mode)
+
     def turn_away_mode_on(self) -> None:
         """Turn away mode on."""
         raise NotImplementedError()

homeassistant/components/water_heater/strings.json

@@ -71,5 +71,13 @@
       "name": "[%key:common::action::turn_off%]",
       "description": "Turns water heater off."
     }
+  },
+  "exceptions": {
+    "not_valid_operation_mode": {
+      "message": "Operation mode {operation_mode} is not valid for {entity_id}. Valid operation modes are: {operation_list}."
+    },
+    "operation_list_not_defined": {
+      "message": "Operation mode {operation_mode} is not valid for {entity_id}. The operation list is not defined."
+    }
   }
 }

tests/components/water_heater/conftest.py

@@ -0,0 +1,22 @@
+"""Fixtures for water heater platform tests."""
+from collections.abc import Generator
+
+import pytest
+
+from homeassistant.config_entries import ConfigFlow
+from homeassistant.core import HomeAssistant
+
+from tests.common import mock_config_flow, mock_platform
+
+
+class MockFlow(ConfigFlow):
+    """Test flow."""
+
+
+@pytest.fixture
+def config_flow_fixture(hass: HomeAssistant) -> Generator[None, None, None]:
+    """Mock config flow."""
+    mock_platform(hass, "test.config_flow")
+
+    with mock_config_flow("test", MockFlow):
+        yield

tests/components/water_heater/test_init.py

@@ -1,6 +1,7 @@
 """The tests for the water heater component."""
 from __future__ import annotations
 
+from unittest import mock
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
@@ -10,17 +11,27 @@ from homeassistant.components import water_heater
 from homeassistant.components.water_heater import (
     ATTR_OPERATION_LIST,
     ATTR_OPERATION_MODE,
+    DOMAIN,
+    SERVICE_SET_OPERATION_MODE,
     SET_TEMPERATURE_SCHEMA,
     WaterHeaterEntity,
     WaterHeaterEntityFeature,
 )
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import UnitOfTemperature
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ServiceValidationError
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from tests.common import (
+    MockConfigEntry,
+    MockModule,
+    MockPlatform,
     async_mock_service,
     help_test_all,
     import_and_test_deprecated_constant_enum,
+    mock_integration,
+    mock_platform,
 )
 
 
@@ -65,9 +76,12 @@ async def test_set_temp_schema(
 class MockWaterHeaterEntity(WaterHeaterEntity):
     """Mock water heater device to use in tests."""
 
-    _attr_operation_list: list[str] = ["off", "heat_pump", "gas"]
+    _attr_operation_list: list[str] | None = ["off", "heat_pump", "gas"]
     _attr_operation = "heat_pump"
     _attr_supported_features = WaterHeaterEntityFeature.ON_OFF
+    _attr_temperature_unit = UnitOfTemperature.CELSIUS
+
+    set_operation_mode: MagicMock = MagicMock()
 
 
 async def test_sync_turn_on(hass: HomeAssistant) -> None:
@@ -106,6 +120,95 @@ async def test_sync_turn_off(hass: HomeAssistant) -> None:
     assert water_heater.async_turn_off.call_count == 1
 
 
+async def test_operation_mode_validation(
+    hass: HomeAssistant, config_flow_fixture: None
+) -> None:
+    """Test operation mode validation."""
+    water_heater_entity = MockWaterHeaterEntity()
+    water_heater_entity.hass = hass
+    water_heater_entity._attr_name = "test"
+    water_heater_entity._attr_unique_id = "test"
+    water_heater_entity._attr_supported_features = (
+        WaterHeaterEntityFeature.OPERATION_MODE
+    )
+    water_heater_entity._attr_current_operation = None
+    water_heater_entity._attr_operation_list = None
+
+    async def async_setup_entry_init(
+        hass: HomeAssistant, config_entry: ConfigEntry
+    ) -> bool:
+        """Set up test config entry."""
+        await hass.config_entries.async_forward_entry_setups(config_entry, [DOMAIN])
+        return True
+
+    async def async_setup_entry_water_heater_platform(
+        hass: HomeAssistant,
+        config_entry: ConfigEntry,
+        async_add_entities: AddEntitiesCallback,
+    ) -> None:
+        """Set up test water_heater platform via config entry."""
+        async_add_entities([water_heater_entity])
+
+    mock_integration(
+        hass,
+        MockModule(
+            "test",
+            async_setup_entry=async_setup_entry_init,
+        ),
+        built_in=False,
+    )
+    mock_platform(
+        hass,
+        "test.water_heater",
+        MockPlatform(async_setup_entry=async_setup_entry_water_heater_platform),
+    )
+
+    config_entry = MockConfigEntry(domain="test")
+    config_entry.add_to_hass(hass)
+    assert await hass.config_entries.async_setup(config_entry.entry_id)
+
+    data = {"entity_id": "water_heater.test", "operation_mode": "test"}
+
+    with pytest.raises(ServiceValidationError) as exc:
+        await hass.services.async_call(
+            DOMAIN, SERVICE_SET_OPERATION_MODE, data, blocking=True
+        )
+    assert (
+        str(exc.value) == "Operation mode test not valid for entity water_heater.test. "
+        "The operation list is not defined"
+    )
+    assert exc.value.translation_domain == DOMAIN
+    assert exc.value.translation_key == "operation_list_not_defined"
+    assert exc.value.translation_placeholders == {
+        "entity_id": "water_heater.test",
+        "operation_mode": "test",
+    }
+
+    water_heater_entity._attr_operation_list = ["gas", "eco"]
+    with pytest.raises(ServiceValidationError) as exc:
+        await hass.services.async_call(
+            DOMAIN, SERVICE_SET_OPERATION_MODE, data, blocking=True
+        )
+    assert (
+        str(exc.value) == "Operation mode test not valid for entity water_heater.test. "
+        "Valid operation modes are: gas, eco"
+    )
+    assert exc.value.translation_domain == DOMAIN
+    assert exc.value.translation_key == "not_valid_operation_mode"
+    assert exc.value.translation_placeholders == {
+        "entity_id": "water_heater.test",
+        "operation_mode": "test",
+        "operation_list": "gas, eco",
+    }
+
+    data = {"entity_id": "water_heater.test", "operation_mode": "eco"}
+    await hass.services.async_call(
+        DOMAIN, SERVICE_SET_OPERATION_MODE, data, blocking=True
+    )
+    await hass.async_block_till_done()
+    water_heater_entity.set_operation_mode.assert_has_calls([mock.call("eco")])
+
+
 def test_all() -> None:
     """Test module.__all__ is correctly set."""
     help_test_all(water_heater)