From 45a927ffb27de1ab8171c5bd18fc2b1b2e666bc2 Mon Sep 17 00:00:00 2001 From: Xiaonan Shen Date: Mon, 31 Aug 2020 22:40:56 +0800 Subject: [PATCH] Add config flow to yeelight (#37191) --- CODEOWNERS | 2 +- .../components/discovery/__init__.py | 2 +- homeassistant/components/yeelight/__init__.py | 330 ++++++++++++++---- .../components/yeelight/binary_sensor.py | 34 +- .../components/yeelight/config_flow.py | 194 ++++++++++ homeassistant/components/yeelight/light.py | 51 ++- .../components/yeelight/manifest.json | 14 +- .../components/yeelight/strings.json | 39 +++ .../components/yeelight/translations/en.json | 39 +++ homeassistant/generated/config_flows.py | 1 + tests/components/yeelight/__init__.py | 17 +- .../components/yeelight/test_binary_sensor.py | 4 +- tests/components/yeelight/test_config_flow.py | 261 ++++++++++++++ tests/components/yeelight/test_init.py | 69 ++++ tests/components/yeelight/test_light.py | 192 +++++----- 15 files changed, 1044 insertions(+), 205 deletions(-) create mode 100644 homeassistant/components/yeelight/config_flow.py create mode 100644 homeassistant/components/yeelight/strings.json create mode 100644 homeassistant/components/yeelight/translations/en.json create mode 100644 tests/components/yeelight/test_config_flow.py create mode 100644 tests/components/yeelight/test_init.py diff --git a/CODEOWNERS b/CODEOWNERS index ecf7745e595..6a2955f3c13 100644 --- a/CODEOWNERS +++ b/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 diff --git a/homeassistant/components/discovery/__init__.py b/homeassistant/components/discovery/__init__.py index 2879d4bfbec..bf3a50c837e 100644 --- a/homeassistant/components/discovery/__init__.py +++ b/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 = ( diff --git a/homeassistant/components/yeelight/__init__.py b/homeassistant/components/yeelight/__init__.py index b0413599fe3..e463e5dad3f 100644 --- a/homeassistant/components/yeelight/__init__.py +++ b/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.components.discovery import SERVICE_YEELIGHT -from homeassistant.components.light import DOMAIN as LIGHT_DOMAIN +from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry, ConfigEntryNotReady 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_connect, dispatcher_send -from homeassistant.helpers.event import track_time_interval +from homeassistant.helpers.dispatcher import dispatcher_send +from homeassistant.helpers.entity import Entity +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): - _LOGGER.debug("Adding autodetected %s", info["hostname"]) - - name = "yeelight_{}_{}".format(info["device_type"], info["properties"]["mac"]) - - device_config = DEVICE_SCHEMA({CONF_NAME: name}) - - _setup_device(hass, config, info[CONF_HOST], device_config) - - discovery.listen(hass, SERVICE_YEELIGHT, device_discovered) - - def update(_): - 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, {} + # Import manually configured devices + for ipaddr, device_config in config.get(DOMAIN, {}).get(CONF_DEVICES, {}).items(): + _LOGGER.debug("Importing configured %s", ipaddr) + entry_config = { + CONF_IP_ADDRESS: ipaddr, + **device_config, + } + hass.async_create_task( + hass.config_entries.flow.async_init( + DOMAIN, + context={"source": SOURCE_IMPORT}, + data=entry_config, + ), ) - 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): - devices = hass.data[DATA_YEELIGHT] +async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: + """Set up Yeelight from a config entry.""" - if ipaddr in devices: - return + async def _initialize(ipaddr: str) -> None: + 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.add_job(device.setup) + hass.data[DOMAIN][DATA_CONFIG_ENTRIES][entry.entry_id] = { + 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) - self._bulb_device = Bulb(self.ipaddr, model=config.get(CONF_MODEL)) + unique_id = bulb.capabilities.get("id") + 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): - """Fetch initial device properties.""" - self._update_properties() + async def async_setup(self): + """Set up the device.""" + + 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() diff --git a/homeassistant/components/yeelight/binary_sensor.py b/homeassistant/components/yeelight/binary_sensor.py index 1696ca9bcb2..ae811cd91d3 100644 --- a/homeassistant/components/yeelight/binary_sensor.py +++ b/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): - """Set up the Yeelight sensors.""" - if not discovery_info: - return - - device = hass.data[DATA_YEELIGHT][discovery_info["host"]] - +async def async_setup_entry( + hass: HomeAssistant, config_entry: ConfigEntry, async_add_entities +) -> None: + """Set up Yeelight from a config entry.""" + device = hass.data[DOMAIN][DATA_CONFIG_ENTRIES][config_entry.entry_id][DATA_DEVICE] 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.""" diff --git a/homeassistant/components/yeelight/config_flow.py b/homeassistant/components/yeelight/config_flow.py new file mode 100644 index 00000000000..656680c9c8b --- /dev/null +++ b/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.""" diff --git a/homeassistant/components/yeelight/light.py b/homeassistant/components/yeelight/light.py index 793af380db9..cc580d60700 100644 --- a/homeassistant/components/yeelight/light.py +++ b/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.core import callback +from homeassistant.config_entries import ConfigEntry +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): - """Set up the Yeelight bulbs.""" - - if not discovery_info: - return +async def async_setup_entry( + hass: HomeAssistant, config_entry: ConfigEntry, async_add_entities +) -> None: + """Set up Yeelight from a config entry.""" 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 = ( - discovery_info.get(CONF_NIGHTLIGHT_SWITCH_TYPE) == NIGHTLIGHT_SWITCH_TYPE_LIGHT - ) + nl_switch_light = device.config.get(CONF_NIGHTLIGHT_SWITCH) 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) - setup_services(hass) + async_add_entities(lights, True) + 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.""" diff --git a/homeassistant/components/yeelight/manifest.json b/homeassistant/components/yeelight/manifest.json index 32ccf1c117e..92baced836a 100644 --- a/homeassistant/components/yeelight/manifest.json +++ b/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"], - "after_dependencies": ["discovery"], - "codeowners": ["@rytilahti", "@zewelor"] -} + "requirements": [ + "yeelight==0.5.2" + ], + "codeowners": [ + "@rytilahti", + "@zewelor", + "@shenxn" + ], + "config_flow": true +} \ No newline at end of file diff --git a/homeassistant/components/yeelight/strings.json b/homeassistant/components/yeelight/strings.json new file mode 100644 index 00000000000..cc52b83f080 --- /dev/null +++ b/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" + } + } + } + } +} \ No newline at end of file diff --git a/homeassistant/components/yeelight/translations/en.json b/homeassistant/components/yeelight/translations/en.json new file mode 100644 index 00000000000..68adafdf376 --- /dev/null +++ b/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" +} \ No newline at end of file diff --git a/homeassistant/generated/config_flows.py b/homeassistant/generated/config_flows.py index 62877b614b1..05ce927c773 100644 --- a/homeassistant/generated/config_flows.py +++ b/homeassistant/generated/config_flows.py @@ -208,6 +208,7 @@ FLOWS = [ "wolflink", "xiaomi_aqara", "xiaomi_miio", + "yeelight", "zerproc", "zha", "zwave" diff --git a/tests/components/yeelight/__init__.py b/tests/components/yeelight/__init__.py index 7f1f7d7d236..a2f5b935947 100644 --- a/tests/components/yeelight/__init__.py +++ b/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) diff --git a/tests/components/yeelight/test_binary_sensor.py b/tests/components/yeelight/test_binary_sensor.py index bf20a7ec5b0..b3281168077 100644 --- a/tests/components/yeelight/test_binary_sensor.py +++ b/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() diff --git a/tests/components/yeelight/test_config_flow.py b/tests/components/yeelight/test_config_flow.py new file mode 100644 index 00000000000..7d1e51afbb1 --- /dev/null +++ b/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 diff --git a/tests/components/yeelight/test_init.py b/tests/components/yeelight/test_init.py new file mode 100644 index 00000000000..004efed8deb --- /dev/null +++ b/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 diff --git a/tests/components/yeelight/test_light.py b/tests/components/yeelight/test_light.py index c44c343e51b..aafe45851d3 100644 --- a/tests/components/yeelight/test_light.py +++ b/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): - await async_setup_component(hass, DOMAIN, YAML_CONFIGURATION) + 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() 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): - capabilities = {**CAPABILITIES} - capabilities["id"] = f"yeelight.{unique_id}" - mocked_bulb = _mocked_bulb() - 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_setup(config_entry): + with patch(f"{MODULE}.Bulb", return_value=mocked_bulb): + await hass.config_entries.async_setup(config_entry.entry_id) + await hass.async_block_till_done() async def _async_test( - name, bulb_type, model, target_properties, nightlight_properties=None, - entity_name=None, - entity_id=None, + name=NAME, + entity_id=ENTITY_LIGHT, ): - if entity_id is None: - entity_id = f"light.{name}" + 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: 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" - entity_id = f"light.{name}" + 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: 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", - entity_id="light.ambient_ambilight", + name=f"{NAME} ambilight", + entity_id=f"{ENTITY_LIGHT}_ambilight", ) async def test_effects(hass: HomeAssistant): """Test effects.""" - yaml_configuration = { - DOMAIN: { - CONF_DEVICES: YAML_CONFIGURATION[DOMAIN][CONF_DEVICES], - CONF_CUSTOM_EFFECTS: [ - { - CONF_NAME: "mock_effect", - CONF_FLOW_PARAMS: { - ATTR_COUNT: 3, - ATTR_TRANSITIONS: [ - {YEELIGHT_HSV_TRANSACTION: [300, 50, 500, 50]}, - {YEELIGHT_RGB_TRANSITION: [100, 100, 100, 300, 30]}, - {YEELIGHT_TEMPERATURE_TRANSACTION: [3000, 200, 20]}, - {YEELIGHT_SLEEP_TRANSACTION: [800]}, - ], + assert await async_setup_component( + hass, + DOMAIN, + { + DOMAIN: { + CONF_CUSTOM_EFFECTS: [ + { + CONF_NAME: "mock_effect", + CONF_FLOW_PARAMS: { + ATTR_COUNT: 3, + ATTR_TRANSITIONS: [ + {YEELIGHT_HSV_TRANSACTION: [300, 50, 500, 50]}, + {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): - assert await async_setup_component(hass, DOMAIN, yaml_configuration) + 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(ENTITY_LIGHT).attributes.get(