From c44391f36809cf8eacdc750b06f8db8c70dc08bf Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Tue, 26 Apr 2022 09:10:48 +0200
Subject: [PATCH] Use climate enums in smarttub (#70738)

---
 homeassistant/components/smarttub/climate.py | 38 ++++++++------------
 tests/components/smarttub/test_climate.py    | 21 +++++------
 2 files changed, 24 insertions(+), 35 deletions(-)

diff --git a/homeassistant/components/smarttub/climate.py b/homeassistant/components/smarttub/climate.py
index 683282f0913..8789c33bc9a 100644
--- a/homeassistant/components/smarttub/climate.py
+++ b/homeassistant/components/smarttub/climate.py
@@ -1,13 +1,15 @@
 """Platform for climate integration."""
+from __future__ import annotations
+
 from smarttub import Spa
 
-from homeassistant.components.climate import ClimateEntity, ClimateEntityFeature
+from homeassistant.components.climate import ClimateEntity
 from homeassistant.components.climate.const import (
-    CURRENT_HVAC_HEAT,
-    CURRENT_HVAC_IDLE,
-    HVAC_MODE_HEAT,
     PRESET_ECO,
     PRESET_NONE,
+    ClimateEntityFeature,
+    HVACAction,
+    HVACMode,
 )
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_TEMPERATURE, TEMP_CELSIUS
@@ -29,8 +31,8 @@ PRESET_MODES = {
 HEAT_MODES = {v: k for k, v in PRESET_MODES.items()}
 
 HVAC_ACTIONS = {
-    "OFF": CURRENT_HVAC_IDLE,
-    "ON": CURRENT_HVAC_HEAT,
+    "OFF": HVACAction.IDLE,
+    "ON": HVACAction.HEATING,
 }
 
 
@@ -51,6 +53,10 @@ async def async_setup_entry(
 class SmartTubThermostat(SmartTubEntity, ClimateEntity):
     """The target water temperature for the spa."""
 
+    # SmartTub devices don't seem to have the option of disabling the heater,
+    # so this is always HVACMode.HEAT.
+    _attr_hvac_mode = HVACMode.HEAT
+    _attr_hvac_modes = [HVACMode.HEAT]
     # Only target temperature is supported.
     _attr_supported_features = (
         ClimateEntityFeature.PRESET_MODE | ClimateEntityFeature.TARGET_TEMPERATURE
@@ -66,30 +72,16 @@ class SmartTubThermostat(SmartTubEntity, ClimateEntity):
         return TEMP_CELSIUS
 
     @property
-    def hvac_action(self):
+    def hvac_action(self) -> HVACAction | None:
         """Return the current running hvac operation."""
         return HVAC_ACTIONS.get(self.spa_status.heater)
 
-    @property
-    def hvac_modes(self):
-        """Return the list of available hvac operation modes."""
-        return [HVAC_MODE_HEAT]
-
-    @property
-    def hvac_mode(self):
-        """Return the current hvac mode.
-
-        SmartTub devices don't seem to have the option of disabling the heater,
-        so this is always HVAC_MODE_HEAT.
-        """
-        return HVAC_MODE_HEAT
-
-    async def async_set_hvac_mode(self, hvac_mode: str):
+    async def async_set_hvac_mode(self, hvac_mode: HVACMode):
         """Set new target hvac mode.
 
         As with hvac_mode, we don't really have an option here.
         """
-        if hvac_mode == HVAC_MODE_HEAT:
+        if hvac_mode == HVACMode.HEAT:
             return
         raise NotImplementedError(hvac_mode)
 
diff --git a/tests/components/smarttub/test_climate.py b/tests/components/smarttub/test_climate.py
index 83a223cee98..b290475d245 100644
--- a/tests/components/smarttub/test_climate.py
+++ b/tests/components/smarttub/test_climate.py
@@ -1,5 +1,4 @@
 """Test the SmartTub climate platform."""
-
 import smarttub
 
 from homeassistant.components.climate.const import (
@@ -11,17 +10,15 @@ from homeassistant.components.climate.const import (
     ATTR_MIN_TEMP,
     ATTR_PRESET_MODE,
     ATTR_PRESET_MODES,
-    CURRENT_HVAC_HEAT,
-    CURRENT_HVAC_IDLE,
     DOMAIN as CLIMATE_DOMAIN,
-    HVAC_MODE_HEAT,
     PRESET_ECO,
     PRESET_NONE,
     SERVICE_SET_HVAC_MODE,
     SERVICE_SET_PRESET_MODE,
     SERVICE_SET_TEMPERATURE,
-    SUPPORT_PRESET_MODE,
-    SUPPORT_TARGET_TEMPERATURE,
+    ClimateEntityFeature,
+    HVACAction,
+    HVACMode,
 )
 from homeassistant.components.smarttub.const import DEFAULT_MAX_TEMP, DEFAULT_MIN_TEMP
 from homeassistant.const import (
@@ -40,19 +37,19 @@ async def test_thermostat_update(spa, spa_state, setup_entry, hass):
     state = hass.states.get(entity_id)
     assert state
 
-    assert state.attributes[ATTR_HVAC_ACTION] == CURRENT_HVAC_HEAT
+    assert state.attributes[ATTR_HVAC_ACTION] == HVACAction.HEATING
 
     spa_state.heater = "OFF"
     await trigger_update(hass)
     state = hass.states.get(entity_id)
 
-    assert state.attributes[ATTR_HVAC_ACTION] == CURRENT_HVAC_IDLE
+    assert state.attributes[ATTR_HVAC_ACTION] == HVACAction.IDLE
 
-    assert set(state.attributes[ATTR_HVAC_MODES]) == {HVAC_MODE_HEAT}
-    assert state.state == HVAC_MODE_HEAT
+    assert set(state.attributes[ATTR_HVAC_MODES]) == {HVACMode.HEAT}
+    assert state.state == HVACMode.HEAT
     assert (
         state.attributes[ATTR_SUPPORTED_FEATURES]
-        == SUPPORT_PRESET_MODE | SUPPORT_TARGET_TEMPERATURE
+        == ClimateEntityFeature.PRESET_MODE | ClimateEntityFeature.TARGET_TEMPERATURE
     )
     assert state.attributes[ATTR_CURRENT_TEMPERATURE] == 38
     assert state.attributes[ATTR_TEMPERATURE] == 39
@@ -71,7 +68,7 @@ async def test_thermostat_update(spa, spa_state, setup_entry, hass):
     await hass.services.async_call(
         CLIMATE_DOMAIN,
         SERVICE_SET_HVAC_MODE,
-        {ATTR_ENTITY_ID: entity_id, ATTR_HVAC_MODE: HVAC_MODE_HEAT},
+        {ATTR_ENTITY_ID: entity_id, ATTR_HVAC_MODE: HVACMode.HEAT},
         blocking=True,
     )
     # does nothing