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

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