Restore yeelight workaround for failing to update state after on/off (#57400)

homeassistant/components/yeelight/__init__.py

@@ -36,8 +36,8 @@ from homeassistant.helpers.typing import ConfigType
 
 _LOGGER = logging.getLogger(__name__)
 
-STATE_CHANGE_TIME = 0.25  # seconds
-
+STATE_CHANGE_TIME = 0.40  # seconds
+POWER_STATE_CHANGE_TIME = 1  # seconds
 
 DOMAIN = "yeelight"
 DATA_YEELIGHT = DOMAIN

homeassistant/components/yeelight/light.py

@@ -38,6 +38,7 @@ from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import entity_platform
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
+from homeassistant.helpers.event import async_call_later
 import homeassistant.util.color as color_util
 from homeassistant.util.color import (
     color_temperature_kelvin_to_mired as kelvin_to_mired,
@@ -62,6 +63,7 @@ from . import (
     DATA_DEVICE,
     DATA_UPDATED,
     DOMAIN,
+    POWER_STATE_CHANGE_TIME,
     YEELIGHT_FLOW_TRANSITION_SCHEMA,
     YeelightEntity,
 )
@@ -247,7 +249,7 @@ def _async_cmd(func):
         except BULB_NETWORK_EXCEPTIONS as ex:
             # A network error happened, the bulb is likely offline now
             self.device.async_mark_unavailable()
-            self.async_write_ha_state()
+            self.async_state_changed()
             exc_message = str(ex) or type(ex)
             raise HomeAssistantError(
                 f"Error when calling {func.__name__} for bulb {self.device.name} at {self.device.host}: {exc_message}"
@@ -419,13 +421,22 @@ class YeelightGenericLight(YeelightEntity, LightEntity):
         else:
             self._custom_effects = {}
 
+        self._unexpected_state_check = None
+
+    @callback
+    def async_state_changed(self):
+        """Call when the device changes state."""
+        if not self._device.available:
+            self._async_cancel_pending_state_check()
+        self.async_write_ha_state()
+
     async def async_added_to_hass(self):
         """Handle entity which will be added."""
         self.async_on_remove(
             async_dispatcher_connect(
                 self.hass,
                 DATA_UPDATED.format(self._device.host),
-                self.async_write_ha_state,
+                self.async_state_changed,
             )
         )
         await super().async_added_to_hass()
@@ -760,6 +771,33 @@ class YeelightGenericLight(YeelightEntity, LightEntity):
         if self.config[CONF_SAVE_ON_CHANGE] and (brightness or colortemp or rgb):
             await self.async_set_default()
 
+        self._async_schedule_state_check(True)
+
+    @callback
+    def _async_cancel_pending_state_check(self):
+        """Cancel a pending state check."""
+        if self._unexpected_state_check:
+            self._unexpected_state_check()
+            self._unexpected_state_check = None
+
+    @callback
+    def _async_schedule_state_check(self, expected_power_state):
+        """Schedule a poll if the change failed to get pushed back to us.
+
+        Some devices (mainly nightlights) will not send back the on state
+        so we need to force a refresh.
+        """
+        self._async_cancel_pending_state_check()
+
+        async def _async_update_if_state_unexpected(*_):
+            self._unexpected_state_check = None
+            if self.is_on != expected_power_state:
+                await self.device.async_update(True)
+
+        self._unexpected_state_check = async_call_later(
+            self.hass, POWER_STATE_CHANGE_TIME, _async_update_if_state_unexpected
+        )
+
     @_async_cmd
     async def _async_turn_off(self, duration) -> None:
         """Turn off with a given transition duration wrapped with _async_cmd."""
@@ -775,6 +813,7 @@ class YeelightGenericLight(YeelightEntity, LightEntity):
             duration = int(kwargs.get(ATTR_TRANSITION) * 1000)  # kwarg in s
 
         await self._async_turn_off(duration)
+        self._async_schedule_state_check(False)
 
     @_async_cmd
     async def async_set_mode(self, mode: str):

tests/components/yeelight/test_light.py

@@ -1,5 +1,6 @@
 """Test the Yeelight light."""
 import asyncio
+from datetime import timedelta
 import logging
 import socket
 from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
@@ -98,6 +99,7 @@ from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import entity_registry as er
 from homeassistant.setup import async_setup_component
+from homeassistant.util import dt as dt_util
 from homeassistant.util.color import (
     color_hs_to_RGB,
     color_hs_to_xy,
@@ -121,7 +123,7 @@ from . import (
     _patch_discovery_interval,
 )
 
-from tests.common import MockConfigEntry
+from tests.common import MockConfigEntry, async_fire_time_changed
 
 CONFIG_ENTRY_DATA = {
     CONF_HOST: IP_ADDRESS,
@@ -1377,3 +1379,73 @@ async def test_ambilight_with_nightlight_disabled(hass: HomeAssistant):
     assert state.state == "on"
     # bg_power off should not set the brightness to 0
     assert state.attributes[ATTR_BRIGHTNESS] == 128
+
+
+async def test_state_fails_to_update_triggers_update(hass: HomeAssistant):
+    """Ensure we call async_get_properties if the turn on/off fails to update the state."""
+    mocked_bulb = _mocked_bulb()
+    properties = {**PROPERTIES}
+    properties.pop("active_mode")
+    properties["color_mode"] = "3"  # HSV
+    mocked_bulb.last_properties = properties
+    mocked_bulb.bulb_type = BulbType.Color
+    config_entry = MockConfigEntry(
+        domain=DOMAIN, data={**CONFIG_ENTRY_DATA, CONF_NIGHTLIGHT_SWITCH: False}
+    )
+    config_entry.add_to_hass(hass)
+    with _patch_discovery(), _patch_discovery_interval(), patch(
+        f"{MODULE}.AsyncBulb", return_value=mocked_bulb
+    ):
+        assert await hass.config_entries.async_setup(config_entry.entry_id)
+        await hass.async_block_till_done()
+        # We use asyncio.create_task now to avoid
+        # blocking starting so we need to block again
+        await hass.async_block_till_done()
+
+    assert len(mocked_bulb.async_get_properties.mock_calls) == 1
+
+    mocked_bulb.last_properties["power"] = "off"
+    await hass.services.async_call(
+        "light",
+        SERVICE_TURN_ON,
+        {
+            ATTR_ENTITY_ID: ENTITY_LIGHT,
+        },
+        blocking=True,
+    )
+    assert len(mocked_bulb.async_turn_on.mock_calls) == 1
+    async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=1))
+    await hass.async_block_till_done()
+    assert len(mocked_bulb.async_get_properties.mock_calls) == 2
+
+    mocked_bulb.last_properties["power"] = "on"
+    for _ in range(5):
+        await hass.services.async_call(
+            "light",
+            SERVICE_TURN_OFF,
+            {
+                ATTR_ENTITY_ID: ENTITY_LIGHT,
+            },
+            blocking=True,
+        )
+    assert len(mocked_bulb.async_turn_off.mock_calls) == 5
+    # Even with five calls we only do one state request
+    # since each successive call should cancel the unexpected
+    # state check
+    async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=2))
+    await hass.async_block_till_done()
+    assert len(mocked_bulb.async_get_properties.mock_calls) == 3
+
+    # But if the state is correct no calls
+    await hass.services.async_call(
+        "light",
+        SERVICE_TURN_ON,
+        {
+            ATTR_ENTITY_ID: ENTITY_LIGHT,
+        },
+        blocking=True,
+    )
+    assert len(mocked_bulb.async_turn_on.mock_calls) == 1
+    async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=3))
+    await hass.async_block_till_done()
+    assert len(mocked_bulb.async_get_properties.mock_calls) == 3