Minor fixes to A. O. Smith integration (#107421)

homeassistant/components/aosmith/water_heater.py

@@ -15,6 +15,7 @@ from homeassistant.components.water_heater import (
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import UnitOfTemperature
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import AOSmithData
@@ -35,13 +36,13 @@ MODE_AOSMITH_TO_HA = {
     AOSmithOperationMode.VACATION: STATE_OFF,
 }
 
-# Operation mode to use when exiting away mode
-DEFAULT_OPERATION_MODE = AOSmithOperationMode.HYBRID
-
-DEFAULT_SUPPORT_FLAGS = (
-    WaterHeaterEntityFeature.TARGET_TEMPERATURE
-    | WaterHeaterEntityFeature.OPERATION_MODE
-)
+# Priority list for operation mode to use when exiting away mode
+# Will use the first mode that is supported by the device
+DEFAULT_OPERATION_MODE_PRIORITY = [
+    AOSmithOperationMode.HYBRID,
+    AOSmithOperationMode.HEAT_PUMP,
+    AOSmithOperationMode.ELECTRIC,
+]
 
 
 async def async_setup_entry(
@@ -93,10 +94,16 @@ class AOSmithWaterHeaterEntity(AOSmithStatusEntity, WaterHeaterEntity):
             for supported_mode in self.device.supported_modes
         )
 
-        if supports_vacation_mode:
-            return DEFAULT_SUPPORT_FLAGS | WaterHeaterEntityFeature.AWAY_MODE
+        support_flags = WaterHeaterEntityFeature.TARGET_TEMPERATURE
 
-        return DEFAULT_SUPPORT_FLAGS
+        # Operation mode only supported if there is more than one mode
+        if len(self.operation_list) > 1:
+            support_flags |= WaterHeaterEntityFeature.OPERATION_MODE
+
+        if supports_vacation_mode:
+            support_flags |= WaterHeaterEntityFeature.AWAY_MODE
+
+        return support_flags
 
     @property
     def target_temperature(self) -> float | None:
@@ -120,6 +127,9 @@ class AOSmithWaterHeaterEntity(AOSmithStatusEntity, WaterHeaterEntity):
 
     async def async_set_operation_mode(self, operation_mode: str) -> None:
         """Set new target operation mode."""
+        if operation_mode not in self.operation_list:
+            raise HomeAssistantError("Operation mode not supported")
+
         aosmith_mode = MODE_HA_TO_AOSMITH.get(operation_mode)
         if aosmith_mode is not None:
             await self.client.update_mode(self.junction_id, aosmith_mode)
@@ -142,6 +152,9 @@ class AOSmithWaterHeaterEntity(AOSmithStatusEntity, WaterHeaterEntity):
 
     async def async_turn_away_mode_off(self) -> None:
         """Turn away mode off."""
-        await self.client.update_mode(self.junction_id, DEFAULT_OPERATION_MODE)
+        supported_aosmith_modes = [x.mode for x in self.device.supported_modes]
 
-        await self.coordinator.async_request_refresh()
+        for mode in DEFAULT_OPERATION_MODE_PRIORITY:
+            if mode in supported_aosmith_modes:
+                await self.client.update_mode(self.junction_id, mode)
+                break

tests/components/aosmith/conftest.py

@@ -29,20 +29,10 @@ FIXTURE_USER_INPUT = {
 
 
 def build_device_fixture(
-    mode_pending: bool, setpoint_pending: bool, has_vacation_mode: bool
+    heat_pump: bool, mode_pending: bool, setpoint_pending: bool, has_vacation_mode: bool
 ):
     """Build a fixture for a device."""
     supported_modes: list[SupportedOperationModeInfo] = [
-        SupportedOperationModeInfo(
-            mode=OperationMode.HYBRID,
-            original_name="HYBRID",
-            has_day_selection=False,
-        ),
-        SupportedOperationModeInfo(
-            mode=OperationMode.HEAT_PUMP,
-            original_name="HEAT_PUMP",
-            has_day_selection=False,
-        ),
         SupportedOperationModeInfo(
             mode=OperationMode.ELECTRIC,
             original_name="ELECTRIC",
@@ -50,6 +40,22 @@ def build_device_fixture(
         ),
     ]
 
+    if heat_pump:
+        supported_modes.append(
+            SupportedOperationModeInfo(
+                mode=OperationMode.HYBRID,
+                original_name="HYBRID",
+                has_day_selection=False,
+            )
+        )
+        supported_modes.append(
+            SupportedOperationModeInfo(
+                mode=OperationMode.HEAT_PUMP,
+                original_name="HEAT_PUMP",
+                has_day_selection=False,
+            )
+        )
+
     if has_vacation_mode:
         supported_modes.append(
             SupportedOperationModeInfo(
@@ -59,10 +65,18 @@ def build_device_fixture(
             )
         )
 
+    device_type = (
+        DeviceType.NEXT_GEN_HEAT_PUMP if heat_pump else DeviceType.RE3_CONNECTED
+    )
+
+    current_mode = OperationMode.HEAT_PUMP if heat_pump else OperationMode.ELECTRIC
+
+    model = "HPTS-50 200 202172000" if heat_pump else "EE12-50H55DVF 100,3806368"
+
     return Device(
         brand="aosmith",
-        model="HPTS-50 200 202172000",
-        device_type=DeviceType.NEXT_GEN_HEAT_PUMP,
+        model=model,
+        device_type=device_type,
         dsn="dsn",
         junction_id="junctionId",
         name="My water heater",
@@ -72,7 +86,7 @@ def build_device_fixture(
         status=DeviceStatus(
             firmware_version="2.14",
             is_online=True,
-            current_mode=OperationMode.HEAT_PUMP,
+            current_mode=current_mode,
             mode_change_pending=mode_pending,
             temperature_setpoint=130,
             temperature_setpoint_pending=setpoint_pending,
@@ -121,6 +135,12 @@ def mock_setup_entry() -> Generator[AsyncMock, None, None]:
         yield mock_setup_entry
 
 
+@pytest.fixture
+def get_devices_fixture_heat_pump() -> bool:
+    """Return whether the device in the get_devices fixture should be a heat pump water heater."""
+    return True
+
+
 @pytest.fixture
 def get_devices_fixture_mode_pending() -> bool:
     """Return whether to set mode_pending in the get_devices fixture."""
@@ -141,6 +161,7 @@ def get_devices_fixture_has_vacation_mode() -> bool:
 
 @pytest.fixture
 async def mock_client(
+    get_devices_fixture_heat_pump: bool,
     get_devices_fixture_mode_pending: bool,
     get_devices_fixture_setpoint_pending: bool,
     get_devices_fixture_has_vacation_mode: bool,
@@ -148,9 +169,10 @@ async def mock_client(
     """Return a mocked client."""
     get_devices_fixture = [
         build_device_fixture(
-            get_devices_fixture_mode_pending,
-            get_devices_fixture_setpoint_pending,
-            get_devices_fixture_has_vacation_mode,
+            heat_pump=get_devices_fixture_heat_pump,
+            mode_pending=get_devices_fixture_mode_pending,
+            setpoint_pending=get_devices_fixture_setpoint_pending,
+            has_vacation_mode=get_devices_fixture_has_vacation_mode,
         )
     ]
     get_all_device_info_fixture = load_json_object_fixture(

tests/components/aosmith/snapshots/test_water_heater.ambr

@@ -8,9 +8,9 @@
       'max_temp': 130,
       'min_temp': 95,
       'operation_list': list([
+        'electric',
         'eco',
         'heat_pump',
-        'electric',
       ]),
       'operation_mode': 'heat_pump',
       'supported_features': <WaterHeaterEntityFeature: 7>,
@@ -25,3 +25,23 @@
     'state': 'heat_pump',
   })
 # ---
+# name: test_state_non_heat_pump[False]
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'away_mode': 'off',
+      'current_temperature': None,
+      'friendly_name': 'My water heater',
+      'max_temp': 130,
+      'min_temp': 95,
+      'supported_features': <WaterHeaterEntityFeature: 5>,
+      'target_temp_high': None,
+      'target_temp_low': None,
+      'temperature': 130,
+    }),
+    'context': <ANY>,
+    'entity_id': 'water_heater.my_water_heater',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'electric',
+  })
+# ---

tests/components/aosmith/test_init.py

@@ -50,7 +50,14 @@ async def test_config_entry_not_ready_get_energy_use_data_error(
     """Test the config entry not ready when get_energy_use_data fails."""
     mock_config_entry.add_to_hass(hass)
 
-    get_devices_fixture = [build_device_fixture(False, False, True)]
+    get_devices_fixture = [
+        build_device_fixture(
+            heat_pump=True,
+            mode_pending=False,
+            setpoint_pending=False,
+            has_vacation_mode=True,
+        )
+    ]
 
     with patch(
         "homeassistant.components.aosmith.config_flow.AOSmithAPIClient.get_devices",

tests/components/aosmith/test_water_heater.py

@@ -25,6 +25,7 @@ from homeassistant.const import (
     ATTR_SUPPORTED_FEATURES,
 )
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import entity_registry as er
 
 from tests.common import MockConfigEntry
@@ -53,6 +54,20 @@ async def test_state(
     assert state == snapshot
 
 
+@pytest.mark.parametrize(
+    ("get_devices_fixture_heat_pump"),
+    [
+        False,
+    ],
+)
+async def test_state_non_heat_pump(
+    hass: HomeAssistant, init_integration: MockConfigEntry, snapshot: SnapshotAssertion
+) -> None:
+    """Test the state of the water heater entity for a non heat pump device."""
+    state = hass.states.get("water_heater.my_water_heater")
+    assert state == snapshot
+
+
 @pytest.mark.parametrize(
     ("get_devices_fixture_has_vacation_mode"),
     [False],
@@ -98,6 +113,24 @@ async def test_set_operation_mode(
     mock_client.update_mode.assert_called_once_with("junctionId", aosmith_mode)
 
 
+async def test_unsupported_operation_mode(
+    hass: HomeAssistant,
+    mock_client: MagicMock,
+    init_integration: MockConfigEntry,
+) -> None:
+    """Test setting the operation mode with an unsupported mode."""
+    with pytest.raises(HomeAssistantError):
+        await hass.services.async_call(
+            WATER_HEATER_DOMAIN,
+            SERVICE_SET_OPERATION_MODE,
+            {
+                ATTR_ENTITY_ID: "water_heater.my_water_heater",
+                ATTR_OPERATION_MODE: "unsupported_mode",
+            },
+            blocking=True,
+        )
+
+
 async def test_set_temperature(
     hass: HomeAssistant,
     mock_client: MagicMock,
@@ -115,10 +148,12 @@ async def test_set_temperature(
 
 
 @pytest.mark.parametrize(
-    ("hass_away_mode", "aosmith_mode"),
+    ("get_devices_fixture_heat_pump", "hass_away_mode", "aosmith_mode"),
     [
-        (True, OperationMode.VACATION),
-        (False, OperationMode.HYBRID),
+        (True, True, OperationMode.VACATION),
+        (True, False, OperationMode.HYBRID),
+        (False, True, OperationMode.VACATION),
+        (False, False, OperationMode.ELECTRIC),
     ],
 )
 async def test_away_mode(