Add config flow to yeelight (#37191)

CODEOWNERS

@@ -490,7 +490,7 @@ homeassistant/components/xiaomi_tv/* @simse
 homeassistant/components/xmpp/* @fabaff @flowolf
 homeassistant/components/yamaha_musiccast/* @jalmeroth
 homeassistant/components/yandex_transport/* @rishatik92 @devbis
-homeassistant/components/yeelight/* @rytilahti @zewelor
+homeassistant/components/yeelight/* @rytilahti @zewelor @shenxn
 homeassistant/components/yeelightsunflower/* @lindsaymarkward
 homeassistant/components/yessssms/* @flowolf
 homeassistant/components/yi/* @bachya

homeassistant/components/discovery/__init__.py

@@ -64,7 +64,6 @@ SERVICE_HANDLERS = {
     SERVICE_KONNECTED: ("konnected", None),
     SERVICE_OCTOPRINT: ("octoprint", None),
     SERVICE_FREEBOX: ("freebox", None),
-    SERVICE_YEELIGHT: ("yeelight", None),
     "yamaha": ("media_player", "yamaha"),
     "frontier_silicon": ("media_player", "frontier_silicon"),
     "openhome": ("media_player", "openhome"),
@@ -93,6 +92,7 @@ MIGRATED_SERVICE_HANDLERS = [
     SERVICE_WEMO,
     SERVICE_XIAOMI_GW,
     "volumio",
+    SERVICE_YEELIGHT,
 ]
 
 DEFAULT_ENABLED = (

homeassistant/components/yeelight/__init__.py

@@ -1,27 +1,26 @@
 """Support for Xiaomi Yeelight WiFi color bulb."""
-
+import asyncio
 from datetime import timedelta
 import logging
 from typing import Optional
 
 import voluptuous as vol
-from yeelight import Bulb, BulbException
+from yeelight import Bulb, BulbException, discover_bulbs
 
-from homeassistant.components.binary_sensor import DOMAIN as BINARY_SENSOR_DOMAIN
+from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry, ConfigEntryNotReady
-from homeassistant.components.discovery import SERVICE_YEELIGHT
-from homeassistant.components.light import DOMAIN as LIGHT_DOMAIN
 from homeassistant.const import (
     ATTR_ENTITY_ID,
     CONF_DEVICES,
-    CONF_HOST,
+    CONF_ID,
+    CONF_IP_ADDRESS,
     CONF_NAME,
     CONF_SCAN_INTERVAL,
 )
-from homeassistant.helpers import discovery
+from homeassistant.core import HomeAssistant, callback
 import homeassistant.helpers.config_validation as cv
-from homeassistant.helpers.discovery import load_platform
+from homeassistant.helpers.dispatcher import dispatcher_send
-from homeassistant.helpers.dispatcher import dispatcher_connect, dispatcher_send
+from homeassistant.helpers.entity import Entity
-from homeassistant.helpers.event import track_time_interval
+from homeassistant.helpers.event import async_track_time_interval
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -32,6 +31,9 @@ DEVICE_INITIALIZED = f"{DOMAIN}_device_initialized"
 
 DEFAULT_NAME = "Yeelight"
 DEFAULT_TRANSITION = 350
+DEFAULT_MODE_MUSIC = False
+DEFAULT_SAVE_ON_CHANGE = False
+DEFAULT_NIGHTLIGHT_SWITCH = False
 
 CONF_MODEL = "model"
 CONF_TRANSITION = "transition"
@@ -40,6 +42,14 @@ CONF_MODE_MUSIC = "use_music_mode"
 CONF_FLOW_PARAMS = "flow_params"
 CONF_CUSTOM_EFFECTS = "custom_effects"
 CONF_NIGHTLIGHT_SWITCH_TYPE = "nightlight_switch_type"
+CONF_NIGHTLIGHT_SWITCH = "nightlight_switch"
+CONF_DEVICE = "device"
+
+DATA_CONFIG_ENTRIES = "config_entries"
+DATA_CUSTOM_EFFECTS = "custom_effects"
+DATA_SCAN_INTERVAL = "scan_interval"
+DATA_DEVICE = "device"
+DATA_UNSUB_UPDATE_LISTENER = "unsub_update_listener"
 
 ATTR_COUNT = "count"
 ATTR_ACTION = "action"
@@ -55,6 +65,7 @@ ACTIVE_COLOR_FLOWING = "1"
 NIGHTLIGHT_SWITCH_TYPE_LIGHT = "light"
 
 SCAN_INTERVAL = timedelta(seconds=30)
+DISCOVERY_INTERVAL = timedelta(seconds=60)
 
 YEELIGHT_RGB_TRANSITION = "RGBTransition"
 YEELIGHT_HSV_TRANSACTION = "HSVTransition"
@@ -139,73 +150,221 @@ UPDATE_REQUEST_PROPERTIES = [
     "active_mode",
 ]
 
+PLATFORMS = ["binary_sensor", "light"]
 
-def setup(hass, config):
+
+async def async_setup(hass: HomeAssistant, config: dict) -> bool:
     """Set up the Yeelight bulbs."""
     conf = config.get(DOMAIN, {})
-    yeelight_data = hass.data[DATA_YEELIGHT] = {}
+    hass.data[DOMAIN] = {
+        DATA_CUSTOM_EFFECTS: conf.get(CONF_CUSTOM_EFFECTS, {}),
+        DATA_CONFIG_ENTRIES: {},
+        DATA_SCAN_INTERVAL: conf.get(CONF_SCAN_INTERVAL, SCAN_INTERVAL),
+    }
 
-    def device_discovered(_, info):
+    # Import manually configured devices
-        _LOGGER.debug("Adding autodetected %s", info["hostname"])
+    for ipaddr, device_config in config.get(DOMAIN, {}).get(CONF_DEVICES, {}).items():
-
+        _LOGGER.debug("Importing configured %s", ipaddr)
-        name = "yeelight_{}_{}".format(info["device_type"], info["properties"]["mac"])
+        entry_config = {
-
+            CONF_IP_ADDRESS: ipaddr,
-        device_config = DEVICE_SCHEMA({CONF_NAME: name})
+            **device_config,
-
+        }
-        _setup_device(hass, config, info[CONF_HOST], device_config)
+        hass.async_create_task(
-
+            hass.config_entries.flow.async_init(
-    discovery.listen(hass, SERVICE_YEELIGHT, device_discovered)
+                DOMAIN,
-
+                context={"source": SOURCE_IMPORT},
-    def update(_):
+                data=entry_config,
-        for device in list(yeelight_data.values()):
+            ),
-            device.update()
-
-    track_time_interval(hass, update, conf.get(CONF_SCAN_INTERVAL, SCAN_INTERVAL))
-
-    def load_platforms(ipaddr):
-        platform_config = hass.data[DATA_YEELIGHT][ipaddr].config.copy()
-        platform_config[CONF_HOST] = ipaddr
-        platform_config[CONF_CUSTOM_EFFECTS] = config.get(DOMAIN, {}).get(
-            CONF_CUSTOM_EFFECTS, {}
         )
-        load_platform(hass, LIGHT_DOMAIN, DOMAIN, platform_config, config)
-        load_platform(hass, BINARY_SENSOR_DOMAIN, DOMAIN, platform_config, config)
-
-    dispatcher_connect(hass, DEVICE_INITIALIZED, load_platforms)
-
-    if DOMAIN in config:
-        for ipaddr, device_config in conf[CONF_DEVICES].items():
-            _LOGGER.debug("Adding configured %s", device_config[CONF_NAME])
-            _setup_device(hass, config, ipaddr, device_config)
 
     return True
 
 
-def _setup_device(hass, _, ipaddr, device_config):
+async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
-    devices = hass.data[DATA_YEELIGHT]
+    """Set up Yeelight from a config entry."""
 
-    if ipaddr in devices:
+    async def _initialize(ipaddr: str) -> None:
-        return
+        device = await _async_setup_device(hass, ipaddr, entry.options)
+        hass.data[DOMAIN][DATA_CONFIG_ENTRIES][entry.entry_id][DATA_DEVICE] = device
+        for component in PLATFORMS:
+            hass.async_create_task(
+                hass.config_entries.async_forward_entry_setup(entry, component)
+            )
 
-    device = YeelightDevice(hass, ipaddr, device_config)
+    # Move options from data for imported entries
+    # Initialize options with default values for other entries
+    if not entry.options:
+        hass.config_entries.async_update_entry(
+            entry,
+            data={
+                CONF_IP_ADDRESS: entry.data.get(CONF_IP_ADDRESS),
+                CONF_ID: entry.data.get(CONF_ID),
+            },
+            options={
+                CONF_NAME: entry.data.get(CONF_NAME, ""),
+                CONF_MODEL: entry.data.get(CONF_MODEL, ""),
+                CONF_TRANSITION: entry.data.get(CONF_TRANSITION, DEFAULT_TRANSITION),
+                CONF_MODE_MUSIC: entry.data.get(CONF_MODE_MUSIC, DEFAULT_MODE_MUSIC),
+                CONF_SAVE_ON_CHANGE: entry.data.get(
+                    CONF_SAVE_ON_CHANGE, DEFAULT_SAVE_ON_CHANGE
+                ),
+                CONF_NIGHTLIGHT_SWITCH: entry.data.get(
+                    CONF_NIGHTLIGHT_SWITCH, DEFAULT_NIGHTLIGHT_SWITCH
+                ),
+            },
+        )
 
-    devices[ipaddr] = device
+    hass.data[DOMAIN][DATA_CONFIG_ENTRIES][entry.entry_id] = {
-    hass.add_job(device.setup)
+        DATA_UNSUB_UPDATE_LISTENER: entry.add_update_listener(_async_update_listener)
+    }
+
+    if entry.data.get(CONF_IP_ADDRESS):
+        # manually added device
+        await _initialize(entry.data[CONF_IP_ADDRESS])
+    else:
+        # discovery
+        scanner = YeelightScanner.async_get(hass)
+        scanner.async_register_callback(entry.data[CONF_ID], _initialize)
+
+    return True
+
+
+async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry):
+    """Unload a config entry."""
+    unload_ok = all(
+        await asyncio.gather(
+            *[
+                hass.config_entries.async_forward_entry_unload(entry, component)
+                for component in PLATFORMS
+            ]
+        )
+    )
+
+    if unload_ok:
+        data = hass.data[DOMAIN][DATA_CONFIG_ENTRIES].pop(entry.entry_id)
+        data[DATA_UNSUB_UPDATE_LISTENER]()
+        data[DATA_DEVICE].async_unload()
+        if entry.data[CONF_ID]:
+            # discovery
+            scanner = YeelightScanner.async_get(hass)
+            scanner.async_unregister_callback(entry.data[CONF_ID])
+
+    return unload_ok
+
+
+async def _async_setup_device(
+    hass: HomeAssistant,
+    ipaddr: str,
+    config: dict,
+) -> None:
+    # Set up device
+    bulb = Bulb(ipaddr, model=config.get(CONF_MODEL) or None)
+    capabilities = await hass.async_add_executor_job(bulb.get_capabilities)
+    if capabilities is None:  # timeout
+        _LOGGER.error("Failed to get capabilities from %s", ipaddr)
+        raise ConfigEntryNotReady
+    device = YeelightDevice(hass, ipaddr, config, bulb)
+    await hass.async_add_executor_job(device.update)
+    await device.async_setup()
+    return device
+
+
+async def _async_update_listener(hass: HomeAssistant, entry: ConfigEntry):
+    """Handle options update."""
+    await hass.config_entries.async_reload(entry.entry_id)
+
+
+class YeelightScanner:
+    """Scan for Yeelight devices."""
+
+    _scanner = None
+
+    @classmethod
+    @callback
+    def async_get(cls, hass: HomeAssistant):
+        """Get scanner instance."""
+        if cls._scanner is None:
+            cls._scanner = cls(hass)
+        return cls._scanner
+
+    def __init__(self, hass: HomeAssistant):
+        """Initialize class."""
+        self._hass = hass
+        self._seen = {}
+        self._callbacks = {}
+        self._scan_task = None
+
+    async def _async_scan(self):
+        _LOGGER.debug("Yeelight scanning")
+        # Run 3 times as packets can get lost
+        for _ in range(3):
+            devices = await self._hass.async_add_executor_job(discover_bulbs)
+            for device in devices:
+                unique_id = device["capabilities"]["id"]
+                if unique_id in self._seen:
+                    continue
+                ipaddr = device["ip"]
+                self._seen[unique_id] = ipaddr
+                _LOGGER.debug("Yeelight discovered at %s", ipaddr)
+                if unique_id in self._callbacks:
+                    self._hass.async_create_task(self._callbacks[unique_id](ipaddr))
+                    self._callbacks.pop(unique_id)
+                    if len(self._callbacks) == 0:
+                        self._async_stop_scan()
+
+        await asyncio.sleep(SCAN_INTERVAL.seconds)
+        self._scan_task = self._hass.loop.create_task(self._async_scan())
+
+    @callback
+    def _async_start_scan(self):
+        """Start scanning for Yeelight devices."""
+        _LOGGER.debug("Start scanning")
+        # Use loop directly to avoid home assistant track this task
+        self._scan_task = self._hass.loop.create_task(self._async_scan())
+
+    @callback
+    def _async_stop_scan(self):
+        """Stop scanning."""
+        _LOGGER.debug("Stop scanning")
+        if self._scan_task is not None:
+            self._scan_task.cancel()
+            self._scan_task = None
+
+    @callback
+    def async_register_callback(self, unique_id, callback_func):
+        """Register callback function."""
+        ipaddr = self._seen.get(unique_id)
+        if ipaddr is not None:
+            self._hass.async_add_job(callback_func(ipaddr))
+        else:
+            self._callbacks[unique_id] = callback_func
+            if len(self._callbacks) == 1:
+                self._async_start_scan()
+
+    @callback
+    def async_unregister_callback(self, unique_id):
+        """Unregister callback function."""
+        if unique_id not in self._callbacks:
+            return
+        self._callbacks.pop(unique_id)
+        if len(self._callbacks) == 0:
+            self._async_stop_scan()
 
 
 class YeelightDevice:
     """Represents single Yeelight device."""
 
-    def __init__(self, hass, ipaddr, config):
+    def __init__(self, hass, ipaddr, config, bulb):
         """Initialize device."""
         self._hass = hass
         self._config = config
         self._ipaddr = ipaddr
-        self._name = config.get(CONF_NAME)
+        unique_id = bulb.capabilities.get("id")
-        self._bulb_device = Bulb(self.ipaddr, model=config.get(CONF_MODEL))
+        self._name = config.get(CONF_NAME) or f"yeelight_{bulb.model}_{unique_id}"
+        self._bulb_device = bulb
         self._device_type = None
         self._available = False
-        self._initialized = False
+        self._remove_time_tracker = None
 
     @property
     def bulb(self):
@@ -237,6 +396,11 @@ class YeelightDevice:
         """Return configured/autodetected device model."""
         return self._bulb_device.model
 
+    @property
+    def fw_version(self):
+        """Return the firmware version."""
+        return self._bulb_device.capabilities.get("fw_ver")
+
     @property
     def is_nightlight_supported(self) -> bool:
         """
@@ -319,8 +483,6 @@ class YeelightDevice:
         try:
             self.bulb.get_properties(UPDATE_REQUEST_PROPERTIES)
             self._available = True
-            if not self._initialized:
-                self._initialize_device()
         except BulbException as ex:
             if self._available:  # just inform once
                 _LOGGER.error(
@@ -348,16 +510,56 @@ class YeelightDevice:
                 ex,
             )
 
-    def _initialize_device(self):
-        self._get_capabilities()
-        self._initialized = True
-        dispatcher_send(self._hass, DEVICE_INITIALIZED, self.ipaddr)
-
     def update(self):
         """Update device properties and send data updated signal."""
         self._update_properties()
         dispatcher_send(self._hass, DATA_UPDATED.format(self._ipaddr))
 
-    def setup(self):
+    async def async_setup(self):
-        """Fetch initial device properties."""
+        """Set up the device."""
-        self._update_properties()
+
+        async def _async_update(_):
+            await self._hass.async_add_executor_job(self.update)
+
+        await _async_update(None)
+        self._remove_time_tracker = async_track_time_interval(
+            self._hass, _async_update, self._hass.data[DOMAIN][DATA_SCAN_INTERVAL]
+        )
+
+    @callback
+    def async_unload(self):
+        """Unload the device."""
+        self._remove_time_tracker()
+
+
+class YeelightEntity(Entity):
+    """Represents single Yeelight entity."""
+
+    def __init__(self, device: YeelightDevice):
+        """Initialize the entity."""
+        self._device = device
+
+    @property
+    def device_info(self) -> dict:
+        """Return the device info."""
+        return {
+            "identifiers": {(DOMAIN, self._device.unique_id)},
+            "name": self._device.name,
+            "manufacturer": "Yeelight",
+            "model": self._device.model,
+            "sw_version": self._device.fw_version,
+        }
+
+    @property
+    def available(self) -> bool:
+        """Return if bulb is available."""
+        return self._device.available
+
+    @property
+    def should_poll(self) -> bool:
+        """No polling needed."""
+        return False
+
+    def update(self) -> None:
+        """Update the entity."""
+        self._device.update()

homeassistant/components/yeelight/binary_sensor.py

@@ -3,32 +3,28 @@ import logging
 from typing import Optional
 
 from homeassistant.components.binary_sensor import BinarySensorEntity
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 
-from . import DATA_UPDATED, DATA_YEELIGHT
+from . import DATA_CONFIG_ENTRIES, DATA_DEVICE, DATA_UPDATED, DOMAIN, YeelightEntity
 
 _LOGGER = logging.getLogger(__name__)
 
 
-def setup_platform(hass, config, add_entities, discovery_info=None):
+async def async_setup_entry(
-    """Set up the Yeelight sensors."""
+    hass: HomeAssistant, config_entry: ConfigEntry, async_add_entities
-    if not discovery_info:
+) -> None:
-        return
+    """Set up Yeelight from a config entry."""
-
+    device = hass.data[DOMAIN][DATA_CONFIG_ENTRIES][config_entry.entry_id][DATA_DEVICE]
-    device = hass.data[DATA_YEELIGHT][discovery_info["host"]]
-
     if device.is_nightlight_supported:
         _LOGGER.debug("Adding nightlight mode sensor for %s", device.name)
-        add_entities([YeelightNightlightModeSensor(device)])
+        async_add_entities([YeelightNightlightModeSensor(device)])
 
 
-class YeelightNightlightModeSensor(BinarySensorEntity):
+class YeelightNightlightModeSensor(YeelightEntity, BinarySensorEntity):
     """Representation of a Yeelight nightlight mode sensor."""
 
-    def __init__(self, device):
-        """Initialize nightlight mode sensor."""
-        self._device = device
-
     async def async_added_to_hass(self):
         """Handle entity which will be added."""
         self.async_on_remove(
@@ -49,16 +45,6 @@ class YeelightNightlightModeSensor(BinarySensorEntity):
 
         return None
 
-    @property
-    def available(self) -> bool:
-        """Return if bulb is available."""
-        return self._device.available
-
-    @property
-    def should_poll(self):
-        """No polling needed."""
-        return False
-
     @property
     def name(self):
         """Return the name of the sensor."""

homeassistant/components/yeelight/config_flow.py

@@ -0,0 +1,194 @@
+"""Config flow for Yeelight integration."""
+import logging
+
+import voluptuous as vol
+import yeelight
+
+from homeassistant import config_entries, exceptions
+from homeassistant.const import CONF_ID, CONF_IP_ADDRESS, CONF_NAME
+from homeassistant.core import callback
+import homeassistant.helpers.config_validation as cv
+
+from . import (
+    CONF_DEVICE,
+    CONF_MODE_MUSIC,
+    CONF_MODEL,
+    CONF_NIGHTLIGHT_SWITCH,
+    CONF_NIGHTLIGHT_SWITCH_TYPE,
+    CONF_SAVE_ON_CHANGE,
+    CONF_TRANSITION,
+    NIGHTLIGHT_SWITCH_TYPE_LIGHT,
+)
+from . import DOMAIN  # pylint:disable=unused-import
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
+    """Handle a config flow for Yeelight."""
+
+    VERSION = 1
+    CONNECTION_CLASS = config_entries.CONN_CLASS_LOCAL_POLL
+
+    @staticmethod
+    @callback
+    def async_get_options_flow(config_entry):
+        """Return the options flow."""
+        return OptionsFlowHandler(config_entry)
+
+    def __init__(self):
+        """Initialize the config flow."""
+        self._capabilities = None
+        self._discovered_devices = {}
+
+    async def async_step_user(self, user_input=None):
+        """Handle the initial step."""
+        errors = {}
+        if user_input is not None:
+            if user_input.get(CONF_IP_ADDRESS):
+                try:
+                    await self._async_try_connect(user_input[CONF_IP_ADDRESS])
+                    return self.async_create_entry(
+                        title=self._async_default_name(),
+                        data=user_input,
+                    )
+                except CannotConnect:
+                    errors["base"] = "cannot_connect"
+                except AlreadyConfigured:
+                    return self.async_abort(reason="already_configured")
+            else:
+                return await self.async_step_pick_device()
+
+        return self.async_show_form(
+            step_id="user",
+            data_schema=vol.Schema({vol.Optional(CONF_IP_ADDRESS): str}),
+            errors=errors,
+        )
+
+    async def async_step_pick_device(self, user_input=None):
+        """Handle the step to pick discovered device."""
+        if user_input is not None:
+            unique_id = user_input[CONF_DEVICE]
+            self._capabilities = self._discovered_devices[unique_id]
+            return self.async_create_entry(
+                title=self._async_default_name(),
+                data={CONF_ID: unique_id},
+            )
+
+        configured_devices = {
+            entry.data[CONF_ID]
+            for entry in self._async_current_entries()
+            if entry.data[CONF_ID]
+        }
+        devices_name = {}
+        # Run 3 times as packets can get lost
+        for _ in range(3):
+            devices = await self.hass.async_add_executor_job(yeelight.discover_bulbs)
+            for device in devices:
+                capabilities = device["capabilities"]
+                unique_id = capabilities["id"]
+                if unique_id in configured_devices:
+                    continue  # ignore configured devices
+                model = capabilities["model"]
+                ipaddr = device["ip"]
+                name = f"{ipaddr} {model} {unique_id}"
+                self._discovered_devices[unique_id] = capabilities
+                devices_name[unique_id] = name
+
+        # Check if there is at least one device
+        if not devices_name:
+            return self.async_abort(reason="no_devices_found")
+        return self.async_show_form(
+            step_id="pick_device",
+            data_schema=vol.Schema({vol.Required(CONF_DEVICE): vol.In(devices_name)}),
+        )
+
+    async def async_step_import(self, user_input=None):
+        """Handle import step."""
+        ipaddr = user_input[CONF_IP_ADDRESS]
+        try:
+            await self._async_try_connect(ipaddr)
+        except CannotConnect:
+            _LOGGER.error("Failed to import %s: cannot connect", ipaddr)
+            return self.async_abort(reason="cannot_connect")
+        except AlreadyConfigured:
+            return self.async_abort(reason="already_configured")
+        if CONF_NIGHTLIGHT_SWITCH_TYPE in user_input:
+            user_input[CONF_NIGHTLIGHT_SWITCH] = (
+                user_input.pop(CONF_NIGHTLIGHT_SWITCH_TYPE)
+                == NIGHTLIGHT_SWITCH_TYPE_LIGHT
+            )
+        return self.async_create_entry(title=user_input[CONF_NAME], data=user_input)
+
+    async def _async_try_connect(self, ipaddr):
+        """Set up with options."""
+        bulb = yeelight.Bulb(ipaddr)
+        try:
+            capabilities = await self.hass.async_add_executor_job(bulb.get_capabilities)
+            if capabilities is None:  # timeout
+                _LOGGER.error("Failed to get capabilities from %s: timeout", ipaddr)
+                raise CannotConnect
+        except OSError as err:
+            _LOGGER.error("Failed to get capabilities from %s: %s", ipaddr, err)
+            raise CannotConnect from err
+        _LOGGER.debug("Get capabilities: %s", capabilities)
+        self._capabilities = capabilities
+        await self.async_set_unique_id(capabilities["id"])
+        self._abort_if_unique_id_configured()
+
+    @callback
+    def _async_default_name(self):
+        model = self._capabilities["model"]
+        unique_id = self._capabilities["id"]
+        return f"yeelight_{model}_{unique_id}"
+
+
+class OptionsFlowHandler(config_entries.OptionsFlow):
+    """Handle a option flow for Yeelight."""
+
+    def __init__(self, config_entry):
+        """Initialize the option flow."""
+        self._config_entry = config_entry
+
+    async def async_step_init(self, user_input=None):
+        """Handle the initial step."""
+        if user_input is not None:
+            # keep the name from imported entries
+            options = {
+                CONF_NAME: self._config_entry.options.get(CONF_NAME),
+                **user_input,
+            }
+            return self.async_create_entry(title="", data=options)
+
+        options = self._config_entry.options
+        return self.async_show_form(
+            step_id="init",
+            data_schema=vol.Schema(
+                {
+                    vol.Optional(CONF_MODEL, default=options[CONF_MODEL]): str,
+                    vol.Required(
+                        CONF_TRANSITION,
+                        default=options[CONF_TRANSITION],
+                    ): cv.positive_int,
+                    vol.Required(
+                        CONF_MODE_MUSIC, default=options[CONF_MODE_MUSIC]
+                    ): bool,
+                    vol.Required(
+                        CONF_SAVE_ON_CHANGE,
+                        default=options[CONF_SAVE_ON_CHANGE],
+                    ): bool,
+                    vol.Required(
+                        CONF_NIGHTLIGHT_SWITCH,
+                        default=options[CONF_NIGHTLIGHT_SWITCH],
+                    ): bool,
+                }
+            ),
+        )
+
+
+class CannotConnect(exceptions.HomeAssistantError):
+    """Error to indicate we cannot connect."""
+
+
+class AlreadyConfigured(exceptions.HomeAssistantError):
+    """Indicate the ip address is already configured."""

homeassistant/components/yeelight/light.py

@@ -1,4 +1,5 @@
 """Light platform support for yeelight."""
+from functools import partial
 import logging
 from typing import Optional
 
@@ -32,8 +33,9 @@ from homeassistant.components.light import (
     SUPPORT_TRANSITION,
     LightEntity,
 )
-from homeassistant.const import ATTR_ENTITY_ID, ATTR_MODE, CONF_HOST, CONF_NAME
+from homeassistant.config_entries import ConfigEntry
-from homeassistant.core import callback
+from homeassistant.const import ATTR_ENTITY_ID, ATTR_MODE, CONF_NAME
+from homeassistant.core import HomeAssistant, callback
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.service import extract_entity_ids
@@ -48,18 +50,20 @@ from . import (
     ATTR_ACTION,
     ATTR_COUNT,
     ATTR_TRANSITIONS,
-    CONF_CUSTOM_EFFECTS,
     CONF_FLOW_PARAMS,
     CONF_MODE_MUSIC,
-    CONF_NIGHTLIGHT_SWITCH_TYPE,
+    CONF_NIGHTLIGHT_SWITCH,
     CONF_SAVE_ON_CHANGE,
     CONF_TRANSITION,
+    DATA_CONFIG_ENTRIES,
+    DATA_CUSTOM_EFFECTS,
+    DATA_DEVICE,
     DATA_UPDATED,
     DATA_YEELIGHT,
     DOMAIN,
-    NIGHTLIGHT_SWITCH_TYPE_LIGHT,
     YEELIGHT_FLOW_TRANSITION_SCHEMA,
     YEELIGHT_SERVICE_SCHEMA,
+    YeelightEntity,
 )
 
 _LOGGER = logging.getLogger(__name__)
@@ -236,22 +240,20 @@ def _cmd(func):
     return _wrap
 
 
-def setup_platform(hass, config, add_entities, discovery_info=None):
+async def async_setup_entry(
-    """Set up the Yeelight bulbs."""
+    hass: HomeAssistant, config_entry: ConfigEntry, async_add_entities
-
+) -> None:
-    if not discovery_info:
+    """Set up Yeelight from a config entry."""
-        return
 
     if PLATFORM_DATA_KEY not in hass.data:
         hass.data[PLATFORM_DATA_KEY] = []
 
-    device = hass.data[DATA_YEELIGHT][discovery_info[CONF_HOST]]
+    custom_effects = _parse_custom_effects(hass.data[DOMAIN][DATA_CUSTOM_EFFECTS])
+
+    device = hass.data[DOMAIN][DATA_CONFIG_ENTRIES][config_entry.entry_id][DATA_DEVICE]
     _LOGGER.debug("Adding %s", device.name)
 
-    custom_effects = _parse_custom_effects(discovery_info[CONF_CUSTOM_EFFECTS])
+    nl_switch_light = device.config.get(CONF_NIGHTLIGHT_SWITCH)
-    nl_switch_light = (
-        discovery_info.get(CONF_NIGHTLIGHT_SWITCH_TYPE) == NIGHTLIGHT_SWITCH_TYPE_LIGHT
-    )
 
     lights = []
 
@@ -290,8 +292,8 @@ def setup_platform(hass, config, add_entities, discovery_info=None):
         )
 
     hass.data[PLATFORM_DATA_KEY] += lights
-    add_entities(lights, True)
+    async_add_entities(lights, True)
-    setup_services(hass)
+    await hass.async_add_executor_job(partial(setup_services, hass))
 
 
 def setup_services(hass):
@@ -406,13 +408,14 @@ def setup_services(hass):
     )
 
 
-class YeelightGenericLight(LightEntity):
+class YeelightGenericLight(YeelightEntity, LightEntity):
     """Representation of a Yeelight generic light."""
 
     def __init__(self, device, custom_effects=None):
         """Initialize the Yeelight light."""
+        super().__init__(device)
+
         self.config = device.config
-        self._device = device
 
         self._brightness = None
         self._color_temp = None
@@ -444,22 +447,12 @@ class YeelightGenericLight(LightEntity):
             )
         )
 
-    @property
-    def should_poll(self):
-        """No polling needed."""
-        return False
-
     @property
     def unique_id(self) -> Optional[str]:
         """Return a unique ID."""
 
         return self.device.unique_id
 
-    @property
-    def available(self) -> bool:
-        """Return if bulb is available."""
-        return self.device.available
-
     @property
     def supported_features(self) -> int:
         """Flag supported features."""

homeassistant/components/yeelight/manifest.json

@@ -2,7 +2,13 @@
   "domain": "yeelight",
   "name": "Yeelight",
   "documentation": "https://www.home-assistant.io/integrations/yeelight",
-  "requirements": ["yeelight==0.5.2"],
+  "requirements": [
-  "after_dependencies": ["discovery"],
+    "yeelight==0.5.2"
-  "codeowners": ["@rytilahti", "@zewelor"]
+  ],
-}
+  "codeowners": [
+    "@rytilahti",
+    "@zewelor",
+    "@shenxn"
+  ],
+  "config_flow": true
+}

homeassistant/components/yeelight/strings.json

@@ -0,0 +1,39 @@
+{
+  "title": "Yeelight",
+  "config": {
+    "step": {
+      "user": {
+        "description": "If you leave IP address empty, discovery will be used to find devices.",
+        "data": {
+          "ip_address": "[%key:common::config_flow::data::ip%]"
+        }
+      },
+      "pick_device": {
+        "data": {
+          "device": "Device"
+        }
+      }
+    },
+    "error": {
+      "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]"
+    },
+    "abort": {
+      "already_configured": "[%key:common::config_flow::abort::already_configured_device%]",
+      "no_devices_found": "[%key:common::config_flow::abort::no_devices_found%]"
+    }
+  },
+  "options": {
+    "step": {
+      "init": {
+        "description": "If you leave model empty, it will be automatically detected.",
+        "data": {
+          "model": "Model (Optional)",
+          "transition": "Transition Time (ms)",
+          "use_music_mode": "Enable Music Mode",
+          "save_on_change": "Save Status On Change",
+          "nightlight_switch": "Use Nightlight Switch"
+        }
+      }
+    }
+  }
+}

homeassistant/components/yeelight/translations/en.json

@@ -0,0 +1,39 @@
+{
+    "config": {
+        "abort": {
+            "already_configured": "[%key:common::config_flow::abort::already_configured_device%]",
+            "no_devices_found": "[%key:common::config_flow::abort::no_devices_found%]"
+        },
+        "error": {
+            "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]"
+        },
+        "step": {
+            "pick_device": {
+                "data": {
+                    "device": "Device"
+                }
+            },
+            "user": {
+                "data": {
+                    "ip_address": "[%key:common::config_flow::data::ip%]"
+                },
+                "description": "If you leave IP address empty, discovery will be used to find devices."
+            }
+        }
+    },
+    "options": {
+        "step": {
+            "init": {
+                "data": {
+                    "model": "Model (Optional)",
+                    "nightlight_switch": "Use Nightlight Switch",
+                    "save_on_change": "Save Status On Change",
+                    "transition": "Transition Time (ms)",
+                    "use_music_mode": "Enable Music Mode"
+                },
+                "description": "If you leave model empty, it will be automatically detected."
+            }
+        }
+    },
+    "title": "Yeelight"
+}

homeassistant/generated/config_flows.py

@@ -208,6 +208,7 @@ FLOWS = [
     "wolflink",
     "xiaomi_aqara",
     "xiaomi_miio",
+    "yeelight",
     "zerproc",
     "zha",
     "zwave"

tests/components/yeelight/__init__.py

@@ -9,9 +9,9 @@ from homeassistant.components.yeelight import (
     DOMAIN,
     NIGHTLIGHT_SWITCH_TYPE_LIGHT,
 )
-from homeassistant.const import CONF_DEVICES, CONF_NAME
+from homeassistant.const import CONF_DEVICES, CONF_ID, CONF_NAME
 
-from tests.async_mock import MagicMock
+from tests.async_mock import MagicMock, patch
 
 IP_ADDRESS = "192.168.1.239"
 MODEL = "color"
@@ -70,6 +70,10 @@ YAML_CONFIGURATION = {
     }
 }
 
+CONFIG_ENTRY_DATA = {
+    CONF_ID: ID,
+}
+
 
 def _mocked_bulb(cannot_connect=False):
     bulb = MagicMock()
@@ -85,3 +89,12 @@ def _mocked_bulb(cannot_connect=False):
     bulb.music_mode = False
 
     return bulb
+
+
+def _patch_discovery(prefix, no_device=False):
+    def _mocked_discovery(timeout=2, interface=False):
+        if no_device:
+            return []
+        return [{"ip": IP_ADDRESS, "port": 55443, "capabilities": CAPABILITIES}]
+
+    return patch(f"{prefix}.discover_bulbs", side_effect=_mocked_discovery)

tests/components/yeelight/test_binary_sensor.py

@@ -12,7 +12,9 @@ from tests.async_mock import patch
 async def test_nightlight(hass: HomeAssistant):
     """Test nightlight sensor."""
     mocked_bulb = _mocked_bulb()
-    with patch(f"{MODULE}.Bulb", return_value=mocked_bulb):
+    with patch(f"{MODULE}.Bulb", return_value=mocked_bulb), patch(
+        f"{MODULE}.config_flow.yeelight.Bulb", return_value=mocked_bulb
+    ):
         await async_setup_component(hass, DOMAIN, YAML_CONFIGURATION)
         await hass.async_block_till_done()
 

tests/components/yeelight/test_config_flow.py

@@ -0,0 +1,261 @@
+"""Test the Yeelight config flow."""
+from homeassistant import config_entries
+from homeassistant.components.yeelight import (
+    CONF_DEVICE,
+    CONF_MODE_MUSIC,
+    CONF_MODEL,
+    CONF_NIGHTLIGHT_SWITCH,
+    CONF_NIGHTLIGHT_SWITCH_TYPE,
+    CONF_SAVE_ON_CHANGE,
+    CONF_TRANSITION,
+    DEFAULT_MODE_MUSIC,
+    DEFAULT_NAME,
+    DEFAULT_NIGHTLIGHT_SWITCH,
+    DEFAULT_SAVE_ON_CHANGE,
+    DEFAULT_TRANSITION,
+    DOMAIN,
+    NIGHTLIGHT_SWITCH_TYPE_LIGHT,
+)
+from homeassistant.const import CONF_ID, CONF_IP_ADDRESS, CONF_NAME
+from homeassistant.core import HomeAssistant
+
+from . import (
+    ID,
+    IP_ADDRESS,
+    MODULE,
+    MODULE_CONFIG_FLOW,
+    NAME,
+    _mocked_bulb,
+    _patch_discovery,
+)
+
+from tests.async_mock import MagicMock, patch
+from tests.common import MockConfigEntry
+
+DEFAULT_CONFIG = {
+    CONF_NAME: NAME,
+    CONF_MODEL: "",
+    CONF_TRANSITION: DEFAULT_TRANSITION,
+    CONF_MODE_MUSIC: DEFAULT_MODE_MUSIC,
+    CONF_SAVE_ON_CHANGE: DEFAULT_SAVE_ON_CHANGE,
+    CONF_NIGHTLIGHT_SWITCH: DEFAULT_NIGHTLIGHT_SWITCH,
+}
+
+
+async def test_discovery(hass: HomeAssistant):
+    """Test setting up discovery."""
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+    assert result["type"] == "form"
+    assert result["step_id"] == "user"
+    assert not result["errors"]
+
+    with _patch_discovery(f"{MODULE_CONFIG_FLOW}.yeelight"):
+        result2 = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {},
+        )
+    assert result2["type"] == "form"
+    assert result2["step_id"] == "pick_device"
+    assert not result2["errors"]
+
+    with patch(f"{MODULE}.async_setup", return_value=True) as mock_setup, patch(
+        f"{MODULE}.async_setup_entry",
+        return_value=True,
+    ) as mock_setup_entry:
+        result3 = await hass.config_entries.flow.async_configure(
+            result["flow_id"], {CONF_DEVICE: ID}
+        )
+
+    assert result3["type"] == "create_entry"
+    assert result3["title"] == NAME
+    assert result3["data"] == {CONF_ID: ID}
+    await hass.async_block_till_done()
+    mock_setup.assert_called_once()
+    mock_setup_entry.assert_called_once()
+
+    # ignore configured devices
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+    assert result["type"] == "form"
+    assert result["step_id"] == "user"
+    assert not result["errors"]
+
+    with _patch_discovery(f"{MODULE_CONFIG_FLOW}.yeelight"):
+        result2 = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {},
+        )
+    assert result2["type"] == "abort"
+    assert result2["reason"] == "no_devices_found"
+
+
+async def test_discovery_no_device(hass: HomeAssistant):
+    """Test discovery without device."""
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+
+    with _patch_discovery(f"{MODULE_CONFIG_FLOW}.yeelight", no_device=True):
+        result2 = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {},
+        )
+
+    assert result2["type"] == "abort"
+    assert result2["reason"] == "no_devices_found"
+
+
+async def test_import(hass: HomeAssistant):
+    """Test import from yaml."""
+    config = {
+        CONF_NAME: DEFAULT_NAME,
+        CONF_IP_ADDRESS: IP_ADDRESS,
+        CONF_TRANSITION: DEFAULT_TRANSITION,
+        CONF_MODE_MUSIC: DEFAULT_MODE_MUSIC,
+        CONF_SAVE_ON_CHANGE: DEFAULT_SAVE_ON_CHANGE,
+        CONF_NIGHTLIGHT_SWITCH_TYPE: NIGHTLIGHT_SWITCH_TYPE_LIGHT,
+    }
+
+    # Cannot connect
+    mocked_bulb = _mocked_bulb(cannot_connect=True)
+    with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb):
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN, context={"source": config_entries.SOURCE_IMPORT}, data=config
+        )
+    type(mocked_bulb).get_capabilities.assert_called_once()
+    assert result["type"] == "abort"
+    assert result["reason"] == "cannot_connect"
+
+    # Success
+    mocked_bulb = _mocked_bulb()
+    with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb), patch(
+        f"{MODULE}.async_setup", return_value=True
+    ) as mock_setup, patch(
+        f"{MODULE}.async_setup_entry",
+        return_value=True,
+    ) as mock_setup_entry:
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN, context={"source": config_entries.SOURCE_IMPORT}, data=config
+        )
+    type(mocked_bulb).get_capabilities.assert_called_once()
+    assert result["type"] == "create_entry"
+    assert result["title"] == DEFAULT_NAME
+    assert result["data"] == {
+        CONF_NAME: DEFAULT_NAME,
+        CONF_IP_ADDRESS: IP_ADDRESS,
+        CONF_TRANSITION: DEFAULT_TRANSITION,
+        CONF_MODE_MUSIC: DEFAULT_MODE_MUSIC,
+        CONF_SAVE_ON_CHANGE: DEFAULT_SAVE_ON_CHANGE,
+        CONF_NIGHTLIGHT_SWITCH: True,
+    }
+    await hass.async_block_till_done()
+    mock_setup.assert_called_once()
+    mock_setup_entry.assert_called_once()
+
+    # Duplicate
+    mocked_bulb = _mocked_bulb()
+    with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb):
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN, context={"source": config_entries.SOURCE_IMPORT}, data=config
+        )
+    assert result["type"] == "abort"
+    assert result["reason"] == "already_configured"
+
+
+async def test_manual(hass: HomeAssistant):
+    """Test manually setup."""
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+    assert result["type"] == "form"
+    assert result["step_id"] == "user"
+    assert not result["errors"]
+
+    # Cannot connect (timeout)
+    mocked_bulb = _mocked_bulb(cannot_connect=True)
+    with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb):
+        result2 = await hass.config_entries.flow.async_configure(
+            result["flow_id"], {CONF_IP_ADDRESS: IP_ADDRESS}
+        )
+    assert result2["type"] == "form"
+    assert result2["step_id"] == "user"
+    assert result2["errors"] == {"base": "cannot_connect"}
+
+    # Cannot connect (error)
+    type(mocked_bulb).get_capabilities = MagicMock(side_effect=OSError)
+    with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb):
+        result3 = await hass.config_entries.flow.async_configure(
+            result["flow_id"], {CONF_IP_ADDRESS: IP_ADDRESS}
+        )
+    assert result3["errors"] == {"base": "cannot_connect"}
+
+    # Success
+    mocked_bulb = _mocked_bulb()
+    with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb), patch(
+        f"{MODULE}.async_setup", return_value=True
+    ), patch(
+        f"{MODULE}.async_setup_entry",
+        return_value=True,
+    ):
+        result4 = await hass.config_entries.flow.async_configure(
+            result["flow_id"], {CONF_IP_ADDRESS: IP_ADDRESS}
+        )
+    assert result4["type"] == "create_entry"
+    assert result4["data"] == {CONF_IP_ADDRESS: IP_ADDRESS}
+
+    # Duplicate
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+    mocked_bulb = _mocked_bulb()
+    with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb):
+        result2 = await hass.config_entries.flow.async_configure(
+            result["flow_id"], {CONF_IP_ADDRESS: IP_ADDRESS}
+        )
+    assert result2["type"] == "abort"
+    assert result2["reason"] == "already_configured"
+
+
+async def test_options(hass: HomeAssistant):
+    """Test options flow."""
+    config_entry = MockConfigEntry(domain=DOMAIN, data={CONF_IP_ADDRESS: IP_ADDRESS})
+    config_entry.add_to_hass(hass)
+
+    mocked_bulb = _mocked_bulb()
+    with patch(f"{MODULE}.Bulb", return_value=mocked_bulb):
+        assert await hass.config_entries.async_setup(config_entry.entry_id)
+        await hass.async_block_till_done()
+
+    config = {
+        CONF_MODEL: "",
+        CONF_TRANSITION: DEFAULT_TRANSITION,
+        CONF_MODE_MUSIC: DEFAULT_MODE_MUSIC,
+        CONF_SAVE_ON_CHANGE: DEFAULT_SAVE_ON_CHANGE,
+        CONF_NIGHTLIGHT_SWITCH: DEFAULT_NIGHTLIGHT_SWITCH,
+    }
+    assert config_entry.options == {
+        CONF_NAME: "",
+        **config,
+    }
+    assert hass.states.get(f"light.{NAME}_nightlight") is None
+
+    result = await hass.config_entries.options.async_init(config_entry.entry_id)
+    assert result["type"] == "form"
+    assert result["step_id"] == "init"
+
+    config[CONF_NIGHTLIGHT_SWITCH] = True
+    with patch(f"{MODULE}.Bulb", return_value=mocked_bulb):
+        result2 = await hass.config_entries.options.async_configure(
+            result["flow_id"], config
+        )
+        await hass.async_block_till_done()
+    assert result2["type"] == "create_entry"
+    assert result2["data"] == {
+        CONF_NAME: "",
+        **config,
+    }
+    assert result2["data"] == config_entry.options
+    assert hass.states.get(f"light.{NAME}_nightlight") is not None

tests/components/yeelight/test_init.py

@@ -0,0 +1,69 @@
+"""Test Yeelight."""
+from homeassistant.components.yeelight import (
+    CONF_NIGHTLIGHT_SWITCH_TYPE,
+    DOMAIN,
+    NIGHTLIGHT_SWITCH_TYPE_LIGHT,
+)
+from homeassistant.const import CONF_DEVICES, CONF_NAME
+from homeassistant.core import HomeAssistant
+from homeassistant.setup import async_setup_component
+
+from . import (
+    CONFIG_ENTRY_DATA,
+    IP_ADDRESS,
+    MODULE,
+    MODULE_CONFIG_FLOW,
+    NAME,
+    _mocked_bulb,
+    _patch_discovery,
+)
+
+from tests.async_mock import patch
+from tests.common import MockConfigEntry
+
+
+async def test_setup_discovery(hass: HomeAssistant):
+    """Test setting up Yeelight by discovery."""
+    config_entry = MockConfigEntry(domain=DOMAIN, data=CONFIG_ENTRY_DATA)
+    config_entry.add_to_hass(hass)
+
+    mocked_bulb = _mocked_bulb()
+    with _patch_discovery(MODULE), patch(f"{MODULE}.Bulb", return_value=mocked_bulb):
+        assert await hass.config_entries.async_setup(config_entry.entry_id)
+        await hass.async_block_till_done()
+
+    assert hass.states.get(f"binary_sensor.{NAME}_nightlight") is not None
+    assert hass.states.get(f"light.{NAME}") is not None
+
+    # Unload
+    assert await hass.config_entries.async_unload(config_entry.entry_id)
+    assert hass.states.get(f"binary_sensor.{NAME}_nightlight") is None
+    assert hass.states.get(f"light.{NAME}") is None
+
+
+async def test_setup_import(hass: HomeAssistant):
+    """Test import from yaml."""
+    mocked_bulb = _mocked_bulb()
+    name = "yeelight"
+    with patch(f"{MODULE}.Bulb", return_value=mocked_bulb), patch(
+        f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb
+    ):
+        assert await async_setup_component(
+            hass,
+            DOMAIN,
+            {
+                DOMAIN: {
+                    CONF_DEVICES: {
+                        IP_ADDRESS: {
+                            CONF_NAME: name,
+                            CONF_NIGHTLIGHT_SWITCH_TYPE: NIGHTLIGHT_SWITCH_TYPE_LIGHT,
+                        }
+                    }
+                }
+            },
+        )
+        await hass.async_block_till_done()
+
+    assert hass.states.get(f"binary_sensor.{name}_nightlight") is not None
+    assert hass.states.get(f"light.{name}") is not None
+    assert hass.states.get(f"light.{name}_nightlight") is not None

tests/components/yeelight/test_light.py

@@ -34,10 +34,15 @@ from homeassistant.components.yeelight import (
     ATTR_TRANSITIONS,
     CONF_CUSTOM_EFFECTS,
     CONF_FLOW_PARAMS,
-    CONF_NIGHTLIGHT_SWITCH_TYPE,
+    CONF_MODE_MUSIC,
+    CONF_NIGHTLIGHT_SWITCH,
+    CONF_SAVE_ON_CHANGE,
+    CONF_TRANSITION,
+    DEFAULT_MODE_MUSIC,
+    DEFAULT_NIGHTLIGHT_SWITCH,
+    DEFAULT_SAVE_ON_CHANGE,
     DEFAULT_TRANSITION,
     DOMAIN,
-    NIGHTLIGHT_SWITCH_TYPE_LIGHT,
     YEELIGHT_HSV_TRANSACTION,
     YEELIGHT_RGB_TRANSITION,
     YEELIGHT_SLEEP_TRANSACTION,
@@ -66,7 +71,7 @@ from homeassistant.components.yeelight.light import (
     YEELIGHT_MONO_EFFECT_LIST,
     YEELIGHT_TEMP_ONLY_EFFECT_LIST,
 )
-from homeassistant.const import ATTR_ENTITY_ID, CONF_DEVICES, CONF_NAME
+from homeassistant.const import ATTR_ENTITY_ID, CONF_ID, CONF_IP_ADDRESS, CONF_NAME
 from homeassistant.core import HomeAssistant
 from homeassistant.setup import async_setup_component
 from homeassistant.util.color import (
@@ -79,24 +84,38 @@ from homeassistant.util.color import (
 )
 
 from . import (
-    CAPABILITIES,
     ENTITY_LIGHT,
     ENTITY_NIGHTLIGHT,
+    IP_ADDRESS,
     MODULE,
     NAME,
     PROPERTIES,
-    YAML_CONFIGURATION,
     _mocked_bulb,
+    _patch_discovery,
 )
 
 from tests.async_mock import MagicMock, patch
+from tests.common import MockConfigEntry
 
 
 async def test_services(hass: HomeAssistant, caplog):
     """Test Yeelight services."""
+    config_entry = MockConfigEntry(
+        domain=DOMAIN,
+        data={
+            CONF_ID: "",
+            CONF_IP_ADDRESS: IP_ADDRESS,
+            CONF_TRANSITION: DEFAULT_TRANSITION,
+            CONF_MODE_MUSIC: True,
+            CONF_SAVE_ON_CHANGE: True,
+            CONF_NIGHTLIGHT_SWITCH: True,
+        },
+    )
+    config_entry.add_to_hass(hass)
+
     mocked_bulb = _mocked_bulb()
-    with patch(f"{MODULE}.Bulb", return_value=mocked_bulb):
+    with _patch_discovery(MODULE), patch(f"{MODULE}.Bulb", return_value=mocked_bulb):
-        await async_setup_component(hass, DOMAIN, YAML_CONFIGURATION)
+        assert await hass.config_entries.async_setup(config_entry.entry_id)
         await hass.async_block_till_done()
 
     async def _async_test_service(service, data, method, payload=None, domain=DOMAIN):
@@ -264,70 +283,70 @@ async def test_services(hass: HomeAssistant, caplog):
 
 async def test_device_types(hass: HomeAssistant):
     """Test different device types."""
+    mocked_bulb = _mocked_bulb()
     properties = {**PROPERTIES}
     properties.pop("active_mode")
     properties["color_mode"] = "3"
+    mocked_bulb.last_properties = properties
 
-    def _create_mocked_bulb(bulb_type, model, unique_id):
+    async def _async_setup(config_entry):
-        capabilities = {**CAPABILITIES}
+        with patch(f"{MODULE}.Bulb", return_value=mocked_bulb):
-        capabilities["id"] = f"yeelight.{unique_id}"
+            await hass.config_entries.async_setup(config_entry.entry_id)
-        mocked_bulb = _mocked_bulb()
+            await hass.async_block_till_done()
-        mocked_bulb.bulb_type = bulb_type
-        mocked_bulb.last_properties = properties
-        mocked_bulb.capabilities = capabilities
-        model_specs = _MODEL_SPECS.get(model)
-        type(mocked_bulb).get_model_specs = MagicMock(return_value=model_specs)
-        return mocked_bulb
-
-    types = {
-        "default": (None, "mono"),
-        "white": (BulbType.White, "mono"),
-        "color": (BulbType.Color, "color"),
-        "white_temp": (BulbType.WhiteTemp, "ceiling1"),
-        "white_temp_mood": (BulbType.WhiteTempMood, "ceiling4"),
-        "ambient": (BulbType.WhiteTempMood, "ceiling4"),
-    }
-
-    devices = {}
-    mocked_bulbs = []
-    unique_id = 0
-    for name, (bulb_type, model) in types.items():
-        devices[f"{name}.yeelight"] = {CONF_NAME: name}
-        devices[f"{name}_nightlight.yeelight"] = {
-            CONF_NAME: f"{name}_nightlight",
-            CONF_NIGHTLIGHT_SWITCH_TYPE: NIGHTLIGHT_SWITCH_TYPE_LIGHT,
-        }
-        mocked_bulbs.append(_create_mocked_bulb(bulb_type, model, unique_id))
-        mocked_bulbs.append(_create_mocked_bulb(bulb_type, model, unique_id + 1))
-        unique_id += 2
-
-    with patch(f"{MODULE}.Bulb", side_effect=mocked_bulbs):
-        await async_setup_component(hass, DOMAIN, {DOMAIN: {CONF_DEVICES: devices}})
-        await hass.async_block_till_done()
 
     async def _async_test(
-        name,
         bulb_type,
         model,
         target_properties,
         nightlight_properties=None,
-        entity_name=None,
+        name=NAME,
-        entity_id=None,
+        entity_id=ENTITY_LIGHT,
     ):
-        if entity_id is None:
+        config_entry = MockConfigEntry(
-            entity_id = f"light.{name}"
+            domain=DOMAIN,
+            data={
+                CONF_ID: "",
+                CONF_IP_ADDRESS: IP_ADDRESS,
+                CONF_TRANSITION: DEFAULT_TRANSITION,
+                CONF_MODE_MUSIC: DEFAULT_MODE_MUSIC,
+                CONF_SAVE_ON_CHANGE: DEFAULT_SAVE_ON_CHANGE,
+                CONF_NIGHTLIGHT_SWITCH: False,
+            },
+        )
+        config_entry.add_to_hass(hass)
+
+        mocked_bulb.bulb_type = bulb_type
+        model_specs = _MODEL_SPECS.get(model)
+        type(mocked_bulb).get_model_specs = MagicMock(return_value=model_specs)
+        await _async_setup(config_entry)
+
         state = hass.states.get(entity_id)
         assert state.state == "on"
-        target_properties["friendly_name"] = entity_name or name
+        target_properties["friendly_name"] = name
         target_properties["flowing"] = False
         target_properties["night_light"] = True
         assert dict(state.attributes) == target_properties
 
+        await hass.config_entries.async_unload(config_entry.entry_id)
+        await config_entry.async_remove(hass)
+
         # nightlight
         if nightlight_properties is None:
             return
-        name += "_nightlight"
+        config_entry = MockConfigEntry(
-        entity_id = f"light.{name}"
+            domain=DOMAIN,
+            data={
+                CONF_ID: "",
+                CONF_IP_ADDRESS: IP_ADDRESS,
+                CONF_TRANSITION: DEFAULT_TRANSITION,
+                CONF_MODE_MUSIC: DEFAULT_MODE_MUSIC,
+                CONF_SAVE_ON_CHANGE: DEFAULT_SAVE_ON_CHANGE,
+                CONF_NIGHTLIGHT_SWITCH: True,
+            },
+        )
+        config_entry.add_to_hass(hass)
+        await _async_setup(config_entry)
+
         assert hass.states.get(entity_id).state == "off"
         state = hass.states.get(f"{entity_id}_nightlight")
         assert state.state == "on"
@@ -337,6 +356,9 @@ async def test_device_types(hass: HomeAssistant):
         nightlight_properties["night_light"] = True
         assert dict(state.attributes) == nightlight_properties
 
+        await hass.config_entries.async_unload(config_entry.entry_id)
+        await config_entry.async_remove(hass)
+
     bright = round(255 * int(PROPERTIES["bright"]) / 100)
     current_brightness = round(255 * int(PROPERTIES["current_brightness"]) / 100)
     ct = color_temperature_kelvin_to_mired(int(PROPERTIES["ct"]))
@@ -355,7 +377,6 @@ async def test_device_types(hass: HomeAssistant):
 
     # Default
     await _async_test(
-        "default",
         None,
         "mono",
         {
@@ -367,7 +388,6 @@ async def test_device_types(hass: HomeAssistant):
 
     # White
     await _async_test(
-        "white",
         BulbType.White,
         "mono",
         {
@@ -380,7 +400,6 @@ async def test_device_types(hass: HomeAssistant):
     # Color
     model_specs = _MODEL_SPECS["color"]
     await _async_test(
-        "color",
         BulbType.Color,
         "color",
         {
@@ -404,7 +423,6 @@ async def test_device_types(hass: HomeAssistant):
     # WhiteTemp
     model_specs = _MODEL_SPECS["ceiling1"]
     await _async_test(
-        "white_temp",
         BulbType.WhiteTemp,
         "ceiling1",
         {
@@ -427,9 +445,10 @@ async def test_device_types(hass: HomeAssistant):
     )
 
     # WhiteTempMood
+    properties.pop("power")
+    properties["main_power"] = "on"
     model_specs = _MODEL_SPECS["ceiling4"]
     await _async_test(
-        "white_temp_mood",
         BulbType.WhiteTempMood,
         "ceiling4",
         {
@@ -454,7 +473,6 @@ async def test_device_types(hass: HomeAssistant):
         },
     )
     await _async_test(
-        "ambient",
         BulbType.WhiteTempMood,
         "ceiling4",
         {
@@ -468,36 +486,52 @@ async def test_device_types(hass: HomeAssistant):
             "rgb_color": bg_rgb_color,
             "xy_color": bg_xy_color,
         },
-        entity_name="ambient ambilight",
+        name=f"{NAME} ambilight",
-        entity_id="light.ambient_ambilight",
+        entity_id=f"{ENTITY_LIGHT}_ambilight",
     )
 
 
 async def test_effects(hass: HomeAssistant):
     """Test effects."""
-    yaml_configuration = {
+    assert await async_setup_component(
-        DOMAIN: {
+        hass,
-            CONF_DEVICES: YAML_CONFIGURATION[DOMAIN][CONF_DEVICES],
+        DOMAIN,
-            CONF_CUSTOM_EFFECTS: [
+        {
-                {
+            DOMAIN: {
-                    CONF_NAME: "mock_effect",
+                CONF_CUSTOM_EFFECTS: [
-                    CONF_FLOW_PARAMS: {
+                    {
-                        ATTR_COUNT: 3,
+                        CONF_NAME: "mock_effect",
-                        ATTR_TRANSITIONS: [
+                        CONF_FLOW_PARAMS: {
-                            {YEELIGHT_HSV_TRANSACTION: [300, 50, 500, 50]},
+                            ATTR_COUNT: 3,
-                            {YEELIGHT_RGB_TRANSITION: [100, 100, 100, 300, 30]},
+                            ATTR_TRANSITIONS: [
-                            {YEELIGHT_TEMPERATURE_TRANSACTION: [3000, 200, 20]},
+                                {YEELIGHT_HSV_TRANSACTION: [300, 50, 500, 50]},
-                            {YEELIGHT_SLEEP_TRANSACTION: [800]},
+                                {YEELIGHT_RGB_TRANSITION: [100, 100, 100, 300, 30]},
-                        ],
+                                {YEELIGHT_TEMPERATURE_TRANSACTION: [3000, 200, 20]},
+                                {YEELIGHT_SLEEP_TRANSACTION: [800]},
+                            ],
+                        },
                     },
-                },
+                ],
-            ],
+            },
-        }
+        },
-    }
+    )
+
+    config_entry = MockConfigEntry(
+        domain=DOMAIN,
+        data={
+            CONF_ID: "",
+            CONF_IP_ADDRESS: IP_ADDRESS,
+            CONF_TRANSITION: DEFAULT_TRANSITION,
+            CONF_MODE_MUSIC: DEFAULT_MODE_MUSIC,
+            CONF_SAVE_ON_CHANGE: DEFAULT_SAVE_ON_CHANGE,
+            CONF_NIGHTLIGHT_SWITCH: DEFAULT_NIGHTLIGHT_SWITCH,
+        },
+    )
+    config_entry.add_to_hass(hass)
 
     mocked_bulb = _mocked_bulb()
-    with patch(f"{MODULE}.Bulb", return_value=mocked_bulb):
+    with _patch_discovery(MODULE), patch(f"{MODULE}.Bulb", return_value=mocked_bulb):
-        assert await async_setup_component(hass, DOMAIN, yaml_configuration)
+        assert await hass.config_entries.async_setup(config_entry.entry_id)
         await hass.async_block_till_done()
 
     assert hass.states.get(ENTITY_LIGHT).attributes.get(