Ensure yeelights resync state if they are busy on first connect (#55333)

This commit is contained in:
J. Nick Koston 2021-08-27 12:43:53 -05:00 committed by GitHub
parent ed19fdd462
commit 10fa63775d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 87 additions and 24 deletions

View File

@ -163,6 +163,8 @@ UPDATE_REQUEST_PROPERTIES = [
"active_mode",
]
BULB_EXCEPTIONS = (BulbException, asyncio.TimeoutError)
PLATFORMS = ["binary_sensor", "light"]
@ -272,7 +274,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
if entry.data.get(CONF_HOST):
try:
device = await _async_get_device(hass, entry.data[CONF_HOST], entry)
except BulbException as ex:
except BULB_EXCEPTIONS 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):
@ -287,7 +289,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
host = urlparse(capabilities["location"]).hostname
try:
await _async_initialize(hass, entry, host)
except BulbException:
except BULB_EXCEPTIONS:
_LOGGER.exception("Failed to connect to bulb at %s", host)
# discovery
@ -552,6 +554,7 @@ class YeelightDevice:
self._device_type = None
self._available = False
self._initialized = False
self._did_first_update = False
self._name = None
@property
@ -647,14 +650,14 @@ class YeelightDevice:
await self.bulb.async_turn_on(
duration=duration, light_type=light_type, power_mode=power_mode
)
except BulbException as ex:
except BULB_EXCEPTIONS as ex:
_LOGGER.error("Unable to turn the bulb on: %s", ex)
async def async_turn_off(self, duration=DEFAULT_TRANSITION, light_type=None):
"""Turn off device."""
try:
await self.bulb.async_turn_off(duration=duration, light_type=light_type)
except BulbException as ex:
except BULB_EXCEPTIONS as ex:
_LOGGER.error(
"Unable to turn the bulb off: %s, %s: %s", self._host, self.name, ex
)
@ -670,7 +673,7 @@ class YeelightDevice:
if not self._initialized:
self._initialized = True
async_dispatcher_send(self._hass, DEVICE_INITIALIZED.format(self._host))
except BulbException as ex:
except BULB_EXCEPTIONS as ex:
if self._available: # just inform once
_LOGGER.error(
"Unable to update device %s, %s: %s", self._host, self.name, ex
@ -696,6 +699,7 @@ class YeelightDevice:
async def async_update(self, force=False):
"""Update device properties and send data updated signal."""
self._did_first_update = True
if not force and self._initialized and self._available:
# No need to poll unless force, already connected
return
@ -705,7 +709,20 @@ class YeelightDevice:
@callback
def async_update_callback(self, data):
"""Update push from device."""
was_available = self._available
self._available = data.get(KEY_CONNECTED, True)
if self._did_first_update and not was_available and self._available:
# On reconnect the properties may be out of sync
#
# We need to make sure the DEVICE_INITIALIZED dispatcher is setup
# before we can update on reconnect by checking self._did_first_update
#
# If the device drops the connection right away, we do not want to
# do a property resync via async_update since its about
# to be called when async_setup_entry reaches the end of the
# function
#
asyncio.create_task(self.async_update(True))
async_dispatcher_send(self._hass, DATA_UPDATED.format(self._host))

View File

@ -6,7 +6,7 @@ import math
import voluptuous as vol
import yeelight
from yeelight import Bulb, BulbException, Flow, RGBTransition, SleepTransition, flows
from yeelight import Bulb, Flow, RGBTransition, SleepTransition, flows
from yeelight.enums import BulbType, LightType, PowerMode, SceneClass
from homeassistant.components.light import (
@ -49,6 +49,7 @@ from . import (
ATTR_COUNT,
ATTR_MODE_MUSIC,
ATTR_TRANSITIONS,
BULB_EXCEPTIONS,
CONF_FLOW_PARAMS,
CONF_MODE_MUSIC,
CONF_NIGHTLIGHT_SWITCH,
@ -241,7 +242,7 @@ def _async_cmd(func):
try:
_LOGGER.debug("Calling %s with %s %s", func, args, kwargs)
return await func(self, *args, **kwargs)
except BulbException as ex:
except BULB_EXCEPTIONS as ex:
_LOGGER.error("Error when calling %s: %s", func, ex)
return _async_wrap
@ -678,7 +679,7 @@ class YeelightGenericLight(YeelightEntity, LightEntity):
flow = Flow(count=count, transitions=transitions)
try:
await self._bulb.async_start_flow(flow, light_type=self.light_type)
except BulbException as ex:
except BULB_EXCEPTIONS as ex:
_LOGGER.error("Unable to set flash: %s", ex)
@_async_cmd
@ -709,7 +710,7 @@ class YeelightGenericLight(YeelightEntity, LightEntity):
try:
await self._bulb.async_start_flow(flow, light_type=self.light_type)
self._effect = effect
except BulbException as ex:
except BULB_EXCEPTIONS as ex:
_LOGGER.error("Unable to set effect: %s", ex)
async def async_turn_on(self, **kwargs) -> None:
@ -737,7 +738,7 @@ class YeelightGenericLight(YeelightEntity, LightEntity):
await self.hass.async_add_executor_job(
self.set_music_mode, self.config[CONF_MODE_MUSIC]
)
except BulbException as ex:
except BULB_EXCEPTIONS as ex:
_LOGGER.error(
"Unable to turn on music mode, consider disabling it: %s", ex
)
@ -750,7 +751,7 @@ class YeelightGenericLight(YeelightEntity, LightEntity):
await self.async_set_brightness(brightness, duration)
await self.async_set_flash(flash)
await self.async_set_effect(effect)
except BulbException as ex:
except BULB_EXCEPTIONS as ex:
_LOGGER.error("Unable to set bulb properties: %s", ex)
return
@ -758,7 +759,7 @@ class YeelightGenericLight(YeelightEntity, LightEntity):
if self.config[CONF_SAVE_ON_CHANGE] and (brightness or colortemp or rgb):
try:
await self.async_set_default()
except BulbException as ex:
except BULB_EXCEPTIONS as ex:
_LOGGER.error("Unable to set the defaults: %s", ex)
return
@ -784,7 +785,7 @@ class YeelightGenericLight(YeelightEntity, LightEntity):
"""Set a power mode."""
try:
await self._bulb.async_set_power_mode(PowerMode[mode.upper()])
except BulbException as ex:
except BULB_EXCEPTIONS as ex:
_LOGGER.error("Unable to set the power mode: %s", ex)
async def async_start_flow(self, transitions, count=0, action=ACTION_RECOVER):
@ -795,7 +796,7 @@ class YeelightGenericLight(YeelightEntity, LightEntity):
)
await self._bulb.async_start_flow(flow, light_type=self.light_type)
except BulbException as ex:
except BULB_EXCEPTIONS as ex:
_LOGGER.error("Unable to set effect: %s", ex)
async def async_set_scene(self, scene_class, *args):
@ -806,7 +807,7 @@ class YeelightGenericLight(YeelightEntity, LightEntity):
"""
try:
await self._bulb.async_set_scene(scene_class, *args)
except BulbException as ex:
except BULB_EXCEPTIONS as ex:
_LOGGER.error("Unable to set scene: %s", ex)

View File

@ -90,11 +90,33 @@ YAML_CONFIGURATION = {
CONFIG_ENTRY_DATA = {CONF_ID: ID}
class MockAsyncBulb:
"""A mock for yeelight.aio.AsyncBulb."""
def __init__(self, model, bulb_type, cannot_connect):
"""Init the mock."""
self.model = model
self.bulb_type = bulb_type
self._async_callback = None
self._cannot_connect = cannot_connect
async def async_listen(self, callback):
"""Mock the listener."""
if self._cannot_connect:
raise BulbException
self._async_callback = callback
async def async_stop_listening(self):
"""Drop the listener."""
self._async_callback = None
def set_capabilities(self, capabilities):
"""Mock setting capabilities."""
self.capabilities = capabilities
def _mocked_bulb(cannot_connect=False):
bulb = MagicMock()
type(bulb).async_listen = AsyncMock(
side_effect=BulbException if cannot_connect else None
)
bulb = MockAsyncBulb(MODEL, BulbType.Color, cannot_connect)
type(bulb).async_get_properties = AsyncMock(
side_effect=BulbException if cannot_connect else None
)
@ -102,14 +124,10 @@ def _mocked_bulb(cannot_connect=False):
side_effect=BulbException if cannot_connect else None
)
type(bulb).get_model_specs = MagicMock(return_value=_MODEL_SPECS[MODEL])
bulb.capabilities = CAPABILITIES.copy()
bulb.model = MODEL
bulb.bulb_type = BulbType.Color
bulb.last_properties = PROPERTIES.copy()
bulb.music_mode = False
bulb.async_get_properties = AsyncMock()
bulb.async_stop_listening = AsyncMock()
bulb.async_update = AsyncMock()
bulb.async_turn_on = AsyncMock()
bulb.async_turn_off = AsyncMock()
@ -122,7 +140,7 @@ def _mocked_bulb(cannot_connect=False):
bulb.async_set_power_mode = AsyncMock()
bulb.async_set_scene = AsyncMock()
bulb.async_set_default = AsyncMock()
bulb.start_music = MagicMock()
return bulb

View File

@ -3,6 +3,7 @@ from datetime import timedelta
from unittest.mock import AsyncMock, patch
from yeelight import BulbException, BulbType
from yeelight.aio import KEY_CONNECTED
from homeassistant.components.yeelight import (
CONF_MODEL,
@ -414,3 +415,29 @@ async def test_async_setup_with_missing_id(hass: HomeAssistant):
assert config_entry.state is ConfigEntryState.LOADED
assert config_entry.data[CONF_ID] == ID
async def test_connection_dropped_resyncs_properties(hass: HomeAssistant):
"""Test handling a connection drop results in a property resync."""
config_entry = MockConfigEntry(
domain=DOMAIN,
unique_id=ID,
data={CONF_HOST: "127.0.0.1"},
options={CONF_NAME: "Test name"},
)
config_entry.add_to_hass(hass)
mocked_bulb = _mocked_bulb()
with _patch_discovery(), _patch_discovery_timeout(), _patch_discovery_interval(), patch(
f"{MODULE}.AsyncBulb", return_value=mocked_bulb
):
await hass.config_entries.async_setup(config_entry.entry_id)
await hass.async_block_till_done()
assert len(mocked_bulb.async_get_properties.mock_calls) == 1
mocked_bulb._async_callback({KEY_CONNECTED: False})
await hass.async_block_till_done()
assert len(mocked_bulb.async_get_properties.mock_calls) == 1
mocked_bulb._async_callback({KEY_CONNECTED: True})
await hass.async_block_till_done()
await hass.async_block_till_done()
assert len(mocked_bulb.async_get_properties.mock_calls) == 2