From 753a13d68df70f45f162bb688d2c1550e6cd95e5 Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Sat, 23 Apr 2022 22:59:21 +0200
Subject: [PATCH] Use ColorMode enum in wiz (#70554)

---
 homeassistant/components/wiz/light.py | 34 +++++++++++++--------------
 1 file changed, 16 insertions(+), 18 deletions(-)

diff --git a/homeassistant/components/wiz/light.py b/homeassistant/components/wiz/light.py
index 9b2d7e6fab4..d534b62cf21 100644
--- a/homeassistant/components/wiz/light.py
+++ b/homeassistant/components/wiz/light.py
@@ -13,11 +13,8 @@ from homeassistant.components.light import (
     ATTR_EFFECT,
     ATTR_RGBW_COLOR,
     ATTR_RGBWW_COLOR,
-    COLOR_MODE_BRIGHTNESS,
-    COLOR_MODE_COLOR_TEMP,
-    COLOR_MODE_RGBW,
-    COLOR_MODE_RGBWW,
     SUPPORT_EFFECT,
+    ColorMode,
     LightEntity,
 )
 from homeassistant.config_entries import ConfigEntry
@@ -32,7 +29,7 @@ from .const import DOMAIN
 from .entity import WizToggleEntity
 from .models import WizData
 
-RGB_WHITE_CHANNELS_COLOR_MODE = {1: COLOR_MODE_RGBW, 2: COLOR_MODE_RGBWW}
+RGB_WHITE_CHANNELS_COLOR_MODE = {1: ColorMode.RGBW, 2: ColorMode.RGBWW}
 
 
 def _async_pilot_builder(**kwargs: Any) -> PilotBuilder:
@@ -79,14 +76,15 @@ class WizBulbEntity(WizToggleEntity, LightEntity):
         super().__init__(wiz_data, name)
         bulb_type: BulbType = self._device.bulbtype
         features: Features = bulb_type.features
-        color_modes = set()
+        self._attr_supported_color_modes: set[ColorMode | str] = set()
         if features.color:
-            color_modes.add(RGB_WHITE_CHANNELS_COLOR_MODE[bulb_type.white_channels])
+            self._attr_supported_color_modes.add(
+                RGB_WHITE_CHANNELS_COLOR_MODE[bulb_type.white_channels]
+            )
         if features.color_tmp:
-            color_modes.add(COLOR_MODE_COLOR_TEMP)
-        if not color_modes and features.brightness:
-            color_modes.add(COLOR_MODE_BRIGHTNESS)
-        self._attr_supported_color_modes = color_modes
+            self._attr_supported_color_modes.add(ColorMode.COLOR_TEMP)
+        if not self._attr_supported_color_modes and features.brightness:
+            self._attr_supported_color_modes.add(ColorMode.BRIGHTNESS)
         self._attr_effect_list = wiz_data.scenes
         if bulb_type.bulb_type != BulbClass.DW:
             kelvin = bulb_type.kelvin_range
@@ -104,21 +102,21 @@ class WizBulbEntity(WizToggleEntity, LightEntity):
         assert color_modes is not None
         if (brightness := state.get_brightness()) is not None:
             self._attr_brightness = max(0, min(255, brightness))
-        if COLOR_MODE_COLOR_TEMP in color_modes and (
+        if ColorMode.COLOR_TEMP in color_modes and (
             color_temp := state.get_colortemp()
         ):
-            self._attr_color_mode = COLOR_MODE_COLOR_TEMP
+            self._attr_color_mode = ColorMode.COLOR_TEMP
             self._attr_color_temp = color_temperature_kelvin_to_mired(color_temp)
         elif (
-            COLOR_MODE_RGBWW in color_modes and (rgbww := state.get_rgbww()) is not None
+            ColorMode.RGBWW in color_modes and (rgbww := state.get_rgbww()) is not None
         ):
             self._attr_rgbww_color = rgbww
-            self._attr_color_mode = COLOR_MODE_RGBWW
-        elif COLOR_MODE_RGBW in color_modes and (rgbw := state.get_rgbw()) is not None:
+            self._attr_color_mode = ColorMode.RGBWW
+        elif ColorMode.RGBW in color_modes and (rgbw := state.get_rgbw()) is not None:
             self._attr_rgbw_color = rgbw
-            self._attr_color_mode = COLOR_MODE_RGBW
+            self._attr_color_mode = ColorMode.RGBW
         else:
-            self._attr_color_mode = COLOR_MODE_BRIGHTNESS
+            self._attr_color_mode = ColorMode.BRIGHTNESS
         self._attr_effect = state.get_scene()
         super()._async_update_attrs()