From 10fa63775d71755414a906e44fd1a365a40cbec0 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Fri, 27 Aug 2021 12:43:53 -0500 Subject: [PATCH] Ensure yeelights resync state if they are busy on first connect (#55333) --- homeassistant/components/yeelight/__init__.py | 27 +++++++++++--- homeassistant/components/yeelight/light.py | 21 +++++------ tests/components/yeelight/__init__.py | 36 ++++++++++++++----- tests/components/yeelight/test_init.py | 27 ++++++++++++++ 4 files changed, 87 insertions(+), 24 deletions(-) diff --git a/homeassistant/components/yeelight/__init__.py b/homeassistant/components/yeelight/__init__.py index 8684e331fad..a0deb0fdf21 100644 --- a/homeassistant/components/yeelight/__init__.py +++ b/homeassistant/components/yeelight/__init__.py @@ -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)) diff --git a/homeassistant/components/yeelight/light.py b/homeassistant/components/yeelight/light.py index be876690b06..e0c21f21fc7 100644 --- a/homeassistant/components/yeelight/light.py +++ b/homeassistant/components/yeelight/light.py @@ -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) diff --git a/tests/components/yeelight/__init__.py b/tests/components/yeelight/__init__.py index 06c0243e918..4a862fa13dd 100644 --- a/tests/components/yeelight/__init__.py +++ b/tests/components/yeelight/__init__.py @@ -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 diff --git a/tests/components/yeelight/test_init.py b/tests/components/yeelight/test_init.py index 4414909d8e0..4b3ac8e0e83 100644 --- a/tests/components/yeelight/test_init.py +++ b/tests/components/yeelight/test_init.py @@ -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