From 5d8b35facee60db0a49c14f4e31ab6c8a282b22e Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Tue, 26 Apr 2022 14:55:44 +0200
Subject: [PATCH] Use light enums in yeelight (#70797)

---
 homeassistant/components/yeelight/light.py | 16 ++++++----------
 tests/components/yeelight/test_light.py    |  6 +++++-
 2 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/homeassistant/components/yeelight/light.py b/homeassistant/components/yeelight/light.py
index 53862a3443b..a368490e51e 100644
--- a/homeassistant/components/yeelight/light.py
+++ b/homeassistant/components/yeelight/light.py
@@ -23,11 +23,9 @@ from homeassistant.components.light import (
     ATTR_TRANSITION,
     FLASH_LONG,
     FLASH_SHORT,
-    SUPPORT_EFFECT,
-    SUPPORT_FLASH,
-    SUPPORT_TRANSITION,
     ColorMode,
     LightEntity,
+    LightEntityFeature,
 )
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_ENTITY_ID, ATTR_MODE, CONF_NAME
@@ -68,8 +66,6 @@ from .entity import YeelightEntity
 
 _LOGGER = logging.getLogger(__name__)
 
-SUPPORT_YEELIGHT = SUPPORT_TRANSITION | SUPPORT_FLASH | SUPPORT_EFFECT
-
 ATTR_MINUTES = "minutes"
 
 SERVICE_SET_MODE = "set_mode"
@@ -410,6 +406,11 @@ class YeelightGenericLight(YeelightEntity, LightEntity):
 
     _attr_color_mode = ColorMode.BRIGHTNESS
     _attr_supported_color_modes = {ColorMode.BRIGHTNESS}
+    _attr_supported_features = (
+        LightEntityFeature.TRANSITION
+        | LightEntityFeature.FLASH
+        | LightEntityFeature.EFFECT
+    )
     _attr_should_poll = False
 
     def __init__(self, device, entry, custom_effects=None):
@@ -452,11 +453,6 @@ class YeelightGenericLight(YeelightEntity, LightEntity):
         )
         await super().async_added_to_hass()
 
-    @property
-    def supported_features(self) -> int:
-        """Flag supported features."""
-        return SUPPORT_YEELIGHT
-
     @property
     def effect_list(self):
         """Return the list of supported effects."""
diff --git a/tests/components/yeelight/test_light.py b/tests/components/yeelight/test_light.py
index 2e37daaf9dd..fad39a052d5 100644
--- a/tests/components/yeelight/test_light.py
+++ b/tests/components/yeelight/test_light.py
@@ -35,6 +35,7 @@ from homeassistant.components.light import (
     FLASH_SHORT,
     SERVICE_TURN_OFF,
     SERVICE_TURN_ON,
+    LightEntityFeature,
 )
 from homeassistant.components.yeelight.const import (
     ATTR_COUNT,
@@ -82,7 +83,6 @@ from homeassistant.components.yeelight.light import (
     SERVICE_SET_MODE,
     SERVICE_SET_MUSIC_MODE,
     SERVICE_START_FLOW,
-    SUPPORT_YEELIGHT,
     YEELIGHT_COLOR_EFFECT_LIST,
     YEELIGHT_MONO_EFFECT_LIST,
     YEELIGHT_TEMP_ONLY_EFFECT_LIST,
@@ -133,6 +133,10 @@ CONFIG_ENTRY_DATA = {
     CONF_NIGHTLIGHT_SWITCH: DEFAULT_NIGHTLIGHT_SWITCH,
 }
 
+SUPPORT_YEELIGHT = (
+    LightEntityFeature.TRANSITION | LightEntityFeature.FLASH | LightEntityFeature.EFFECT
+)
+
 
 async def test_services(hass: HomeAssistant, caplog):
     """Test Yeelight services."""