From ae1d2926cfb250c115c5d7514e82d6df85e428c7 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Thu, 26 Aug 2021 13:25:26 -0500
Subject: [PATCH] Fix some yeelights showing wrong state after on/off (#55279)

---
 homeassistant/components/yeelight/__init__.py |  6 +-
 homeassistant/components/yeelight/light.py    |  7 +++
 tests/components/yeelight/test_light.py       | 56 +++++++++++++++++++
 3 files changed, 66 insertions(+), 3 deletions(-)

diff --git a/homeassistant/components/yeelight/__init__.py b/homeassistant/components/yeelight/__init__.py
index cee6cb7b9e1..3d1a6cd03e1 100644
--- a/homeassistant/components/yeelight/__init__.py
+++ b/homeassistant/components/yeelight/__init__.py
@@ -692,10 +692,10 @@ class YeelightDevice:
         else:
             self._name = self._host  # Default name is host
 
-    async def async_update(self):
+    async def async_update(self, force=False):
         """Update device properties and send data updated signal."""
-        if self._initialized and self._available:
-            # No need to poll, already connected
+        if not force and self._initialized and self._available:
+            # No need to poll unless force, already connected
             return
         await self._async_update_properties()
         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 4766d897909..bc4d027ce93 100644
--- a/homeassistant/components/yeelight/light.py
+++ b/homeassistant/components/yeelight/light.py
@@ -762,6 +762,10 @@ class YeelightGenericLight(YeelightEntity, LightEntity):
                 _LOGGER.error("Unable to set the defaults: %s", ex)
                 return
 
+        # Some devices (mainly nightlights) will not send back the on state so we need to force a refresh
+        if not self.is_on:
+            await self.device.async_update(True)
+
     async def async_turn_off(self, **kwargs) -> None:
         """Turn off."""
         if not self.is_on:
@@ -772,6 +776,9 @@ class YeelightGenericLight(YeelightEntity, LightEntity):
             duration = int(kwargs.get(ATTR_TRANSITION) * 1000)  # kwarg in s
 
         await self.device.async_turn_off(duration=duration, light_type=self.light_type)
+        # Some devices will not send back the off state so we need to force a refresh
+        if self.is_on:
+            await self.device.async_update(True)
 
     async def async_set_mode(self, mode: str):
         """Set a power mode."""
diff --git a/tests/components/yeelight/test_light.py b/tests/components/yeelight/test_light.py
index 4b8717f4ba4..17924facfad 100644
--- a/tests/components/yeelight/test_light.py
+++ b/tests/components/yeelight/test_light.py
@@ -1122,3 +1122,59 @@ async def test_effects(hass: HomeAssistant):
     for name, target in effects.items():
         await _async_test_effect(name, target)
     await _async_test_effect("not_existed", called=False)
+
+
+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(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()
+
+    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
+    assert len(mocked_bulb.async_get_properties.mock_calls) == 2
+
+    mocked_bulb.last_properties["power"] = "on"
+    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) == 1
+    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
+    assert len(mocked_bulb.async_get_properties.mock_calls) == 3