evlli/hass-core
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add discovery to yeelight (#50385)

Browse source
This commit is contained in:
J. Nick Koston 2021-05-11 19:55:50 -05:00 • committed by GitHub
parent 4e08d22a74
commit c037ebb27c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
9 changed files with 367 additions and 85 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

168
homeassistant/components/yeelight/__init__.py
View file

@ -8,7 +8,7 @@ import logging
import voluptuous as vol
from yeelight import Bulb, BulbException, discover_bulbs
from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry
from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry, ConfigEntryNotReady
from homeassistant.const import (
CONF_DEVICES,
CONF_HOST,
@ -48,8 +48,8 @@ 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"
DATA_REMOVE_INIT_DISPATCHER = "remove_init_dispatcher"
DATA_PLATFORMS_LOADED = "platforms_loaded"
ATTR_COUNT = "count"
ATTR_ACTION = "action"
@ -179,81 +179,115 @@ async def async_setup(hass: HomeAssistant, config: dict) -> bool:
return True
async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Set up Yeelight from a config entry."""
async def _async_initialize(
hass: HomeAssistant,
entry: ConfigEntry,
host: str,
device: YeelightDevice | None = None,
) -> None:
entry_data = hass.data[DOMAIN][DATA_CONFIG_ENTRIES][entry.entry_id] = {
DATA_PLATFORMS_LOADED: False
}
entry.async_on_unload(entry.add_update_listener(_async_update_listener))
async def _initialize(host: str, capabilities: dict | None = None) -> None:
remove_dispatcher = async_dispatcher_connect(
hass,
DEVICE_INITIALIZED.format(host),
_load_platforms,
)
hass.data[DOMAIN][DATA_CONFIG_ENTRIES][entry.entry_id][
DATA_REMOVE_INIT_DISPATCHER
] = remove_dispatcher
device = await _async_get_device(hass, host, entry, capabilities)
hass.data[DOMAIN][DATA_CONFIG_ENTRIES][entry.entry_id][DATA_DEVICE] = device
await device.async_setup()
async def _load_platforms():
@callback
def _async_load_platforms():
if entry_data[DATA_PLATFORMS_LOADED]:
return
entry_data[DATA_PLATFORMS_LOADED] = True
hass.config_entries.async_setup_platforms(entry, PLATFORMS)
# 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_HOST: entry.data.get(CONF_HOST),
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
),
},
)
if not device:
device = await _async_get_device(hass, host, entry)
entry_data[DATA_DEVICE] = device
hass.data[DOMAIN][DATA_CONFIG_ENTRIES][entry.entry_id] = {
DATA_UNSUB_UPDATE_LISTENER: entry.add_update_listener(_async_update_listener)
}
entry.async_on_unload(
async_dispatcher_connect(
hass,
DEVICE_INITIALIZED.format(host),
_async_load_platforms,
)
)
entry.async_on_unload(device.async_unload)
await device.async_setup()
@callback
def _async_populate_entry_options(hass: HomeAssistant, entry: ConfigEntry) -> None:
"""Move options from data for imported entries.
Initialize options with default values for other entries.
"""
if entry.options:
return
hass.config_entries.async_update_entry(
entry,
data={
CONF_HOST: entry.data.get(CONF_HOST),
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
),
},
)
async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Set up Yeelight from a config entry."""
_async_populate_entry_options(hass, entry)
if entry.data.get(CONF_HOST):
# manually added device
await _initialize(entry.data[CONF_HOST])
else:
# discovery
scanner = YeelightScanner.async_get(hass)
scanner.async_register_callback(entry.data[CONF_ID], _initialize)
try:
device = await _async_get_device(hass, entry.data[CONF_HOST], entry)
except OSError as ex:
# If CONF_ID is not valid we cannot fallback to discovery
# so we must retry by raising ConfigEntryNotReady
if not entry.data.get(CONF_ID):
raise ConfigEntryNotReady from ex
# Otherwise fall through to discovery
else:
# manually added device
await _async_initialize(hass, entry, entry.data[CONF_HOST], device=device)
return True
# discovery
scanner = YeelightScanner.async_get(hass)
async def _async_from_discovery(host: str) -> None:
await _async_initialize(hass, entry, host)
scanner.async_register_callback(entry.data[CONF_ID], _async_from_discovery)
return True
async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry):
"""Unload a config entry."""
unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
if unload_ok:
data = hass.data[DOMAIN][DATA_CONFIG_ENTRIES].pop(entry.entry_id)
remove_init_dispatcher = data.get(DATA_REMOVE_INIT_DISPATCHER)
if remove_init_dispatcher is not None:
remove_init_dispatcher()
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])
data_config_entries = hass.data[DOMAIN][DATA_CONFIG_ENTRIES]
entry_data = data_config_entries[entry.entry_id]
return unload_ok
if entry_data[DATA_PLATFORMS_LOADED]:
if not await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
return False
if entry.data.get(CONF_ID):
# discovery
scanner = YeelightScanner.async_get(hass)
scanner.async_unregister_callback(entry.data[CONF_ID])
data_config_entries.pop(entry.entry_id)
return True
@callback
@ -582,16 +616,12 @@ async def _async_get_device(
hass: HomeAssistant,
host: str,
entry: ConfigEntry,
capabilities: dict | None,
) -> YeelightDevice:
# Get model from config and capabilities
model = entry.options.get(CONF_MODEL)
if not model and capabilities is not None:
model = capabilities.get("model")
# Set up device
bulb = Bulb(host, model=model or None)
if capabilities is None:
capabilities = await hass.async_add_executor_job(bulb.get_capabilities)
capabilities = await hass.async_add_executor_job(bulb.get_capabilities)
return YeelightDevice(hass, host, entry.options, bulb, capabilities)

75
homeassistant/components/yeelight/config_flow.py
View file

@ -5,6 +5,7 @@ import voluptuous as vol
import yeelight
from homeassistant import config_entries, exceptions
from homeassistant.components.dhcp import IP_ADDRESS
from homeassistant.const import CONF_DEVICE, CONF_HOST, CONF_ID, CONF_NAME
from homeassistant.core import callback
import homeassistant.helpers.config_validation as cv
@ -21,6 +22,8 @@ from . import (
_async_unique_name,
)
MODEL_UNKNOWN = "unknown"
_LOGGER = logging.getLogger(__name__)
@ -38,22 +41,69 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
def __init__(self):
"""Initialize the config flow."""
self._discovered_devices = {}
self._discovered_model = None
self._discovered_ip = None
async def async_step_homekit(self, discovery_info):
"""Handle discovery from homekit."""
self._discovered_ip = discovery_info["host"]
return await self._async_handle_discovery()
async def async_step_dhcp(self, discovery_info):
"""Handle discovery from dhcp."""
self._discovered_ip = discovery_info[IP_ADDRESS]
return await self._async_handle_discovery()
async def _async_handle_discovery(self):
"""Handle any discovery."""
self.context[CONF_HOST] = self._discovered_ip
for progress in self._async_in_progress():
if progress.get("context", {}).get(CONF_HOST) == self._discovered_ip:
return self.async_abort(reason="already_in_progress")
self._discovered_model = await self._async_try_connect(self._discovered_ip)
if not self.unique_id:
return self.async_abort(reason="cannot_connect")
self._abort_if_unique_id_configured(
updates={CONF_HOST: self._discovered_ip}, reload_on_update=False
)
return await self.async_step_discovery_confirm()
async def async_step_discovery_confirm(self, user_input=None):
"""Confirm discovery."""
if user_input is not None:
return self.async_create_entry(
title=f"{self._discovered_model} {self.unique_id}",
data={CONF_ID: self.unique_id, CONF_HOST: self._discovered_ip},
)
self._set_confirm_only()
placeholders = {
"model": self._discovered_model,
"host": self._discovered_ip,
}
self.context["title_placeholders"] = placeholders
return self.async_show_form(
step_id="discovery_confirm", description_placeholders=placeholders
)
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_HOST):
try:
await self._async_try_connect(user_input[CONF_HOST])
return self.async_create_entry(
title=user_input[CONF_HOST],
data=user_input,
)
except CannotConnect:
errors["base"] = "cannot_connect"
else:
if not user_input.get(CONF_HOST):
return await self.async_step_pick_device()
try:
model = await self._async_try_connect(user_input[CONF_HOST])
except CannotConnect:
errors["base"] = "cannot_connect"
else:
self._abort_if_unique_id_configured()
return self.async_create_entry(
title=f"{model} {self.unique_id}",
data=user_input,
)
user_input = user_input or {}
return self.async_show_form(
@ -117,6 +167,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
user_input.pop(CONF_NIGHTLIGHT_SWITCH_TYPE)
== NIGHTLIGHT_SWITCH_TYPE_LIGHT
)
self._abort_if_unique_id_configured()
return self.async_create_entry(title=user_input[CONF_NAME], data=user_input)
async def _async_try_connect(self, host):
@ -131,8 +182,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
else:
_LOGGER.debug("Get capabilities: %s", capabilities)
await self.async_set_unique_id(capabilities["id"])
self._abort_if_unique_id_configured()
return
return capabilities["model"]
except OSError as err:
_LOGGER.debug("Failed to get capabilities from %s: %s", host, err)
# Ignore the error since get_capabilities uses UDP discovery packet
@ -145,6 +195,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
_LOGGER.error("Failed to get properties from %s: %s", host, err)
raise CannotConnect from err
_LOGGER.debug("Get properties: %s", bulb.last_properties)
return MODEL_UNKNOWN
class OptionsFlowHandler(config_entries.OptionsFlow):

8
homeassistant/components/yeelight/manifest.json
View file

@ -5,5 +5,11 @@
"requirements": ["yeelight==0.6.2"],
"codeowners": ["@rytilahti", "@zewelor", "@shenxn"],
"config_flow": true,
"iot_class": "local_polling"
"iot_class": "local_polling",
"dhcp": [{
"hostname": "yeelink-*"
}],
"homekit": {
"models": ["YLDP*"]
}
}

4
homeassistant/components/yeelight/strings.json
View file

@ -1,5 +1,6 @@
{
"config": {
"flow_title": "{model} {host}",
"step": {
"user": {
"description": "If you leave the host empty, discovery will be used to find devices.",
@ -11,6 +12,9 @@
"data": {
"device": "Device"
}
},
"discovery_confirm": {
"description": "Do you want to setup {model} ({host})?"
}
},
"error": {

4
homeassistant/components/yeelight/translations/en.json
View file

@ -7,7 +7,11 @@
"error": {
"cannot_connect": "Failed to connect"
},
"flow_title": "{model} {host}",
"step": {
"discovery_confirm": {
"description": "Do you want to setup {model} ({host})?"
},
"pick_device": {
"data": {
"device": "Device"

4
homeassistant/generated/dhcp.py
View file

@ -343,5 +343,9 @@ DHCP = [
{
"domain": "verisure",
"macaddress": "0023C1*"
},
{
"domain": "yeelight",
"hostname": "yeelink-*"
}
]

1
homeassistant/generated/zeroconf.py
View file

@ -240,6 +240,7 @@ HOMEKIT = {
"Touch HD": "rainmachine",
"Welcome": "netatmo",
"Wemo": "wemo",
"YLDP*": "yeelight",
"iSmartGate": "gogogate2",
"iZone": "izone",
"tado": "tado"

107
tests/components/yeelight/test_config_flow.py
View file

@ -1,7 +1,9 @@
"""Test the Yeelight config flow."""
from unittest.mock import MagicMock, patch
from homeassistant import config_entries
import pytest
from homeassistant import config_entries, setup
from homeassistant.components.yeelight import (
CONF_MODE_MUSIC,
CONF_MODEL,
@ -19,6 +21,7 @@ from homeassistant.components.yeelight import (
)
from homeassistant.const import CONF_DEVICE, CONF_HOST, CONF_ID, CONF_NAME
from homeassistant.core import HomeAssistant
from homeassistant.data_entry_flow import RESULT_TYPE_ABORT, RESULT_TYPE_FORM
from . import (
ID,
@ -205,7 +208,7 @@ async def test_manual(hass: HomeAssistant):
)
await hass.async_block_till_done()
assert result4["type"] == "create_entry"
assert result4["title"] == IP_ADDRESS
assert result4["title"] == "color 0x000000000015243f"
assert result4["data"] == {CONF_HOST: IP_ADDRESS}
# Duplicate
@ -286,3 +289,103 @@ async def test_manual_no_capabilities(hass: HomeAssistant):
type(mocked_bulb).get_properties.assert_called_once()
assert result["type"] == "create_entry"
assert result["data"] == {CONF_HOST: IP_ADDRESS}
async def test_discovered_by_homekit_and_dhcp(hass):
"""Test we get the form with homekit and abort for dhcp source when we get both."""
await setup.async_setup_component(hass, "persistent_notification", {})
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_HOMEKIT},
data={"host": "1.2.3.4", "properties": {"id": "aa:bb:cc:dd:ee:ff"}},
)
assert result["type"] == RESULT_TYPE_FORM
assert result["errors"] is None
with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb):
result2 = await hass.config_entries.flow.async_init(
DOMAIN,
context={"source": config_entries.SOURCE_DHCP},
data={"ip": "1.2.3.4", "macaddress": "aa:bb:cc:dd:ee:ff"},
)
assert result2["type"] == RESULT_TYPE_ABORT
assert result2["reason"] == "already_in_progress"
with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb):
result3 = await hass.config_entries.flow.async_init(
DOMAIN,
context={"source": config_entries.SOURCE_DHCP},
data={"ip": "1.2.3.4", "macaddress": "00:00:00:00:00:00"},
)
assert result3["type"] == RESULT_TYPE_ABORT
assert result3["reason"] == "already_in_progress"
@pytest.mark.parametrize(
"source, data",
[
(
config_entries.SOURCE_DHCP,
{"ip": IP_ADDRESS, "macaddress": "aa:bb:cc:dd:ee:ff"},
),
(
config_entries.SOURCE_HOMEKIT,
{"host": IP_ADDRESS, "properties": {"id": "aa:bb:cc:dd:ee:ff"}},
),
],
)
async def test_discovered_by_dhcp_or_homekit(hass, source, data):
"""Test we can setup when discovered from dhcp or homekit."""
await setup.async_setup_component(hass, "persistent_notification", {})
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": source},
data=data,
)
assert result["type"] == RESULT_TYPE_FORM
assert result["errors"] is None
with patch(f"{MODULE}.async_setup", return_value=True) as mock_async_setup, patch(
f"{MODULE}.async_setup_entry",
return_value=True,
) as mock_async_setup_entry:
result2 = await hass.config_entries.flow.async_configure(result["flow_id"], {})
assert result2["type"] == "create_entry"
assert result2["data"] == {CONF_HOST: IP_ADDRESS, CONF_ID: "0x000000000015243f"}
assert mock_async_setup.called
assert mock_async_setup_entry.called
@pytest.mark.parametrize(
"source, data",
[
(
config_entries.SOURCE_DHCP,
{"ip": IP_ADDRESS, "macaddress": "aa:bb:cc:dd:ee:ff"},
),
(
config_entries.SOURCE_HOMEKIT,
{"host": IP_ADDRESS, "properties": {"id": "aa:bb:cc:dd:ee:ff"}},
),
],
)
async def test_discovered_by_dhcp_or_homekit_failed_to_get_id(hass, source, data):
"""Test we abort if we cannot get the unique id when discovered from dhcp or homekit."""
await setup.async_setup_component(hass, "persistent_notification", {})
mocked_bulb = _mocked_bulb()
type(mocked_bulb).get_capabilities = MagicMock(return_value=None)
with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb):
result = await hass.config_entries.flow.async_init(
DOMAIN,
context={"source": source},
data=data,
)
assert result["type"] == RESULT_TYPE_ABORT
assert result["reason"] == "cannot_connect"

81
tests/components/yeelight/test_init.py
View file

@ -11,7 +11,14 @@ from homeassistant.components.yeelight import (
DOMAIN,
NIGHTLIGHT_SWITCH_TYPE_LIGHT,
)
from homeassistant.const import CONF_DEVICES, CONF_HOST, CONF_NAME, STATE_UNAVAILABLE
from homeassistant.config_entries import ENTRY_STATE_SETUP_RETRY
from homeassistant.const import (
CONF_DEVICES,
CONF_HOST,
CONF_ID,
CONF_NAME,
STATE_UNAVAILABLE,
)
from homeassistant.core import HomeAssistant
from homeassistant.helpers import entity_registry as er
from homeassistant.setup import async_setup_component
@ -35,6 +42,77 @@ from . import (
from tests.common import MockConfigEntry
async def test_ip_changes_fallback_discovery(hass: HomeAssistant):
"""Test Yeelight ip changes and we fallback to discovery."""
config_entry = MockConfigEntry(
domain=DOMAIN,
data={
CONF_ID: ID,
CONF_HOST: "5.5.5.5",
},
unique_id=ID,
)
config_entry.add_to_hass(hass)
mocked_bulb = _mocked_bulb(True)
mocked_bulb.bulb_type = BulbType.WhiteTempMood
mocked_bulb.get_capabilities = MagicMock(
side_effect=[OSError, CAPABILITIES, CAPABILITIES]
)
_discovered_devices = [
{
"capabilities": CAPABILITIES,
"ip": IP_ADDRESS,
}
]
with patch(f"{MODULE}.Bulb", return_value=mocked_bulb), patch(
f"{MODULE}.discover_bulbs", return_value=_discovered_devices
):
assert await hass.config_entries.async_setup(config_entry.entry_id)
await hass.async_block_till_done()
binary_sensor_entity_id = ENTITY_BINARY_SENSOR_TEMPLATE.format(
f"yeelight_color_{ID}"
)
entity_registry = er.async_get(hass)
assert entity_registry.async_get(binary_sensor_entity_id) is None
await hass.async_block_till_done()
type(mocked_bulb).get_properties = MagicMock(None)
hass.data[DOMAIN][DATA_CONFIG_ENTRIES][config_entry.entry_id][DATA_DEVICE].update()
await hass.async_block_till_done()
await hass.async_block_till_done()
entity_registry = er.async_get(hass)
assert entity_registry.async_get(binary_sensor_entity_id) is not None
async def test_ip_changes_id_missing_cannot_fallback(hass: HomeAssistant):
"""Test Yeelight ip changes and we fallback to discovery."""
config_entry = MockConfigEntry(
domain=DOMAIN,
data={
CONF_HOST: "5.5.5.5",
},
)
config_entry.add_to_hass(hass)
mocked_bulb = _mocked_bulb(True)
mocked_bulb.bulb_type = BulbType.WhiteTempMood
mocked_bulb.get_capabilities = MagicMock(
side_effect=[OSError, CAPABILITIES, CAPABILITIES]
)
with patch(f"{MODULE}.Bulb", return_value=mocked_bulb):
assert not await hass.config_entries.async_setup(config_entry.entry_id)
await hass.async_block_till_done()
assert config_entry.state == ENTRY_STATE_SETUP_RETRY
async def test_setup_discovery(hass: HomeAssistant):
"""Test setting up Yeelight by discovery."""
config_entry = MockConfigEntry(domain=DOMAIN, data=CONFIG_ENTRY_DATA)
@ -182,6 +260,7 @@ async def test_bulb_off_while_adding_in_ha(hass: HomeAssistant):
hass.data[DOMAIN][DATA_CONFIG_ENTRIES][config_entry.entry_id][DATA_DEVICE].update()
await hass.async_block_till_done()
await hass.async_block_till_done()
entity_registry = er.async_get(hass)
assert entity_registry.async_get(binary_sensor_entity_id) is not None