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

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))
 
 

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)
 
 

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()
+    bulb = MockAsyncBulb(MODEL, BulbType.Color, cannot_connect)
-    type(bulb).async_listen = AsyncMock(
-        side_effect=BulbException if cannot_connect else None
-    )
     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
 
 

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