From c104383ab90ede7dd9c92a844b9ddbb4bafa7bf1 Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Tue, 26 Apr 2022 09:49:50 +0200
Subject: [PATCH] Use climate enums in zhong_hong (#70755)

---
 .../components/zhong_hong/climate.py          | 49 +++++++------------
 1 file changed, 18 insertions(+), 31 deletions(-)

diff --git a/homeassistant/components/zhong_hong/climate.py b/homeassistant/components/zhong_hong/climate.py
index 46fbce9db33..633dc50bcb7 100644
--- a/homeassistant/components/zhong_hong/climate.py
+++ b/homeassistant/components/zhong_hong/climate.py
@@ -7,18 +7,11 @@ import voluptuous as vol
 from zhong_hong_hvac.hub import ZhongHongGateway
 from zhong_hong_hvac.hvac import HVAC as ZhongHongHVAC
 
-from homeassistant.components.climate import (
-    PLATFORM_SCHEMA,
-    ClimateEntity,
-    ClimateEntityFeature,
-)
+from homeassistant.components.climate import PLATFORM_SCHEMA, ClimateEntity
 from homeassistant.components.climate.const import (
     ATTR_HVAC_MODE,
-    HVAC_MODE_COOL,
-    HVAC_MODE_DRY,
-    HVAC_MODE_FAN_ONLY,
-    HVAC_MODE_HEAT,
-    HVAC_MODE_OFF,
+    ClimateEntityFeature,
+    HVACMode,
 )
 from homeassistant.const import (
     ATTR_TEMPERATURE,
@@ -56,14 +49,6 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
     }
 )
 
-SUPPORT_HVAC = [
-    HVAC_MODE_COOL,
-    HVAC_MODE_HEAT,
-    HVAC_MODE_DRY,
-    HVAC_MODE_FAN_ONLY,
-    HVAC_MODE_OFF,
-]
-
 ZHONG_HONG_MODE_COOL = "cool"
 ZHONG_HONG_MODE_HEAT = "heat"
 ZHONG_HONG_MODE_DRY = "dry"
@@ -71,10 +56,10 @@ ZHONG_HONG_MODE_FAN_ONLY = "fan_only"
 
 
 MODE_TO_STATE = {
-    ZHONG_HONG_MODE_COOL: HVAC_MODE_COOL,
-    ZHONG_HONG_MODE_HEAT: HVAC_MODE_HEAT,
-    ZHONG_HONG_MODE_DRY: HVAC_MODE_DRY,
-    ZHONG_HONG_MODE_FAN_ONLY: HVAC_MODE_FAN_ONLY,
+    ZHONG_HONG_MODE_COOL: HVACMode.COOL,
+    ZHONG_HONG_MODE_HEAT: HVACMode.HEAT,
+    ZHONG_HONG_MODE_DRY: HVACMode.DRY,
+    ZHONG_HONG_MODE_FAN_ONLY: HVACMode.FAN_ONLY,
 }
 
 
@@ -132,6 +117,13 @@ def setup_platform(
 class ZhongHongClimate(ClimateEntity):
     """Representation of a ZhongHong controller support HVAC."""
 
+    _attr_hvac_modes = [
+        HVACMode.COOL,
+        HVACMode.HEAT,
+        HVACMode.DRY,
+        HVACMode.FAN_ONLY,
+        HVACMode.OFF,
+    ]
     _attr_supported_features = (
         ClimateEntityFeature.TARGET_TEMPERATURE | ClimateEntityFeature.FAN_MODE
     )
@@ -189,16 +181,11 @@ class ZhongHongClimate(ClimateEntity):
         return TEMP_CELSIUS
 
     @property
-    def hvac_mode(self):
+    def hvac_mode(self) -> HVACMode:
         """Return current operation ie. heat, cool, idle."""
         if self.is_on:
             return self._current_operation
-        return HVAC_MODE_OFF
-
-    @property
-    def hvac_modes(self):
-        """Return the list of available operation modes."""
-        return SUPPORT_HVAC
+        return HVACMode.OFF
 
     @property
     def current_temperature(self):
@@ -256,9 +243,9 @@ class ZhongHongClimate(ClimateEntity):
         if (operation_mode := kwargs.get(ATTR_HVAC_MODE)) is not None:
             self.set_hvac_mode(operation_mode)
 
-    def set_hvac_mode(self, hvac_mode):
+    def set_hvac_mode(self, hvac_mode: HVACMode) -> None:
         """Set new target operation mode."""
-        if hvac_mode == HVAC_MODE_OFF:
+        if hvac_mode == HVACMode.OFF:
             if self.is_on:
                 self.turn_off()
             return