From 6157dfe68b19a41a1ff39d657dd8cf08ff857b9d Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Mon, 13 Dec 2021 23:16:54 +0100
Subject: [PATCH] Use new HumidifierDeviceClass enum in generic_hygrostat
 (#61607)

Co-authored-by: epenet <epenet@users.noreply.github.com>
---
 .../components/generic_hygrostat/__init__.py  |  7 ++----
 .../generic_hygrostat/humidifier.py           | 24 ++++++++++++-------
 2 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/homeassistant/components/generic_hygrostat/__init__.py b/homeassistant/components/generic_hygrostat/__init__.py
index b58c98b1d0d..20877f63369 100644
--- a/homeassistant/components/generic_hygrostat/__init__.py
+++ b/homeassistant/components/generic_hygrostat/__init__.py
@@ -2,10 +2,7 @@
 
 import voluptuous as vol
 
-from homeassistant.components.humidifier.const import (
-    DEVICE_CLASS_DEHUMIDIFIER,
-    DEVICE_CLASS_HUMIDIFIER,
-)
+from homeassistant.components.humidifier import HumidifierDeviceClass
 from homeassistant.const import CONF_NAME
 from homeassistant.helpers import config_validation as cv, discovery
 
@@ -34,7 +31,7 @@ HYGROSTAT_SCHEMA = vol.Schema(
         vol.Required(CONF_HUMIDIFIER): cv.entity_id,
         vol.Required(CONF_SENSOR): cv.entity_id,
         vol.Optional(CONF_DEVICE_CLASS): vol.In(
-            [DEVICE_CLASS_HUMIDIFIER, DEVICE_CLASS_DEHUMIDIFIER]
+            [HumidifierDeviceClass.HUMIDIFIER, HumidifierDeviceClass.DEHUMIDIFIER]
         ),
         vol.Optional(CONF_MAX_HUMIDITY): vol.Coerce(int),
         vol.Optional(CONF_MIN_DUR): vol.All(cv.time_period, cv.positive_timedelta),
diff --git a/homeassistant/components/generic_hygrostat/humidifier.py b/homeassistant/components/generic_hygrostat/humidifier.py
index 383674a7f75..7a92eb79d51 100644
--- a/homeassistant/components/generic_hygrostat/humidifier.py
+++ b/homeassistant/components/generic_hygrostat/humidifier.py
@@ -2,11 +2,13 @@
 import asyncio
 import logging
 
-from homeassistant.components.humidifier import PLATFORM_SCHEMA, HumidifierEntity
+from homeassistant.components.humidifier import (
+    PLATFORM_SCHEMA,
+    HumidifierDeviceClass,
+    HumidifierEntity,
+)
 from homeassistant.components.humidifier.const import (
     ATTR_HUMIDITY,
-    DEVICE_CLASS_DEHUMIDIFIER,
-    DEVICE_CLASS_HUMIDIFIER,
     MODE_AWAY,
     MODE_NORMAL,
     SUPPORT_MODES,
@@ -146,7 +148,7 @@ class GenericHygrostat(HumidifierEntity, RestoreEntity):
         self._remove_stale_tracking = None
         self._is_away = False
         if not self._device_class:
-            self._device_class = DEVICE_CLASS_HUMIDIFIER
+            self._device_class = HumidifierDeviceClass.HUMIDIFIER
 
     async def async_added_to_hass(self):
         """Run when entity about to be added."""
@@ -185,7 +187,7 @@ class GenericHygrostat(HumidifierEntity, RestoreEntity):
             if old_state.state:
                 self._state = old_state.state == STATE_ON
         if self._target_humidity is None:
-            if self._device_class == DEVICE_CLASS_HUMIDIFIER:
+            if self._device_class == HumidifierDeviceClass.HUMIDIFIER:
                 self._target_humidity = self.min_humidity
             else:
                 self._target_humidity = self.max_humidity
@@ -396,8 +398,10 @@ class GenericHygrostat(HumidifierEntity, RestoreEntity):
             too_dry = self._target_humidity - self._cur_humidity >= dry_tolerance
             too_wet = self._cur_humidity - self._target_humidity >= wet_tolerance
             if self._is_device_active:
-                if (self._device_class == DEVICE_CLASS_HUMIDIFIER and too_wet) or (
-                    self._device_class == DEVICE_CLASS_DEHUMIDIFIER and too_dry
+                if (
+                    self._device_class == HumidifierDeviceClass.HUMIDIFIER and too_wet
+                ) or (
+                    self._device_class == HumidifierDeviceClass.DEHUMIDIFIER and too_dry
                 ):
                     _LOGGER.info("Turning off humidifier %s", self._switch_entity_id)
                     await self._async_device_turn_off()
@@ -405,8 +409,10 @@ class GenericHygrostat(HumidifierEntity, RestoreEntity):
                     # The time argument is passed only in keep-alive case
                     await self._async_device_turn_on()
             else:
-                if (self._device_class == DEVICE_CLASS_HUMIDIFIER and too_dry) or (
-                    self._device_class == DEVICE_CLASS_DEHUMIDIFIER and too_wet
+                if (
+                    self._device_class == HumidifierDeviceClass.HUMIDIFIER and too_dry
+                ) or (
+                    self._device_class == HumidifierDeviceClass.DEHUMIDIFIER and too_wet
                 ):
                     _LOGGER.info("Turning on humidifier %s", self._switch_entity_id)
                     await self._async_device_turn_on()