From 2ac075bb3762e93e92be7b0bc7ab96185559a64b Mon Sep 17 00:00:00 2001
From: Eric Severance <esev@esev.com>
Date: Wed, 17 Feb 2021 15:38:08 -0800
Subject: [PATCH] Perform wemo state update quickly after a timeout (#46702)

* Perform state update quickly after a timeout

* with async_timeout.timeout(...) -> async with async_timeout.timeout(...)
---
 homeassistant/components/wemo/entity.py      | 19 ++++++++++++----
 tests/components/wemo/entity_test_helpers.py | 24 +++++++++++++-------
 tests/components/wemo/test_light_bridge.py   |  7 +++---
 3 files changed, 35 insertions(+), 15 deletions(-)

diff --git a/homeassistant/components/wemo/entity.py b/homeassistant/components/wemo/entity.py
index 91470d0cd5c..d9d90c5508b 100644
--- a/homeassistant/components/wemo/entity.py
+++ b/homeassistant/components/wemo/entity.py
@@ -89,16 +89,28 @@ class WemoEntity(Entity):
             return
 
         try:
-            with async_timeout.timeout(self.platform.scan_interval.seconds - 0.1):
-                await asyncio.shield(self._async_locked_update(True))
+            async with async_timeout.timeout(
+                self.platform.scan_interval.seconds - 0.1
+            ) as timeout:
+                await asyncio.shield(self._async_locked_update(True, timeout))
         except asyncio.TimeoutError:
             _LOGGER.warning("Lost connection to %s", self.name)
             self._available = False
 
-    async def _async_locked_update(self, force_update: bool) -> None:
+    async def _async_locked_update(
+        self, force_update: bool, timeout: Optional[async_timeout.timeout] = None
+    ) -> None:
         """Try updating within an async lock."""
         async with self._update_lock:
             await self.hass.async_add_executor_job(self._update, force_update)
+            # When the timeout expires HomeAssistant is no longer waiting for an
+            # update from the device. Instead, the state needs to be updated
+            # asynchronously. This also handles the case where an update came
+            # directly from the device (device push). In that case no polling
+            # update was involved and the state also needs to be updated
+            # asynchronously.
+            if not timeout or timeout.expired:
+                self.async_write_ha_state()
 
 
 class WemoSubscriptionEntity(WemoEntity):
@@ -152,4 +164,3 @@ class WemoSubscriptionEntity(WemoEntity):
             return
 
         await self._async_locked_update(force_update)
-        self.async_write_ha_state()
diff --git a/tests/components/wemo/entity_test_helpers.py b/tests/components/wemo/entity_test_helpers.py
index 4c92640572b..e584cb5fb39 100644
--- a/tests/components/wemo/entity_test_helpers.py
+++ b/tests/components/wemo/entity_test_helpers.py
@@ -6,6 +6,7 @@ import asyncio
 import threading
 from unittest.mock import patch
 
+import async_timeout
 from pywemo.ouimeaux_device.api.service import ActionException
 
 from homeassistant.components.homeassistant import (
@@ -146,7 +147,19 @@ async def test_async_update_with_timeout_and_recovery(hass, wemo_entity, pywemo_
     assert hass.states.get(wemo_entity.entity_id).state == STATE_OFF
     await async_setup_component(hass, HA_DOMAIN, {})
 
-    with patch("async_timeout.timeout", side_effect=asyncio.TimeoutError):
+    event = threading.Event()
+
+    def get_state(*args):
+        event.wait()
+        return 0
+
+    if hasattr(pywemo_device, "bridge_update"):
+        pywemo_device.bridge_update.side_effect = get_state
+    else:
+        pywemo_device.get_state.side_effect = get_state
+    timeout = async_timeout.timeout(0)
+
+    with patch("async_timeout.timeout", return_value=timeout):
         await hass.services.async_call(
             HA_DOMAIN,
             SERVICE_UPDATE_ENTITY,
@@ -157,11 +170,6 @@ async def test_async_update_with_timeout_and_recovery(hass, wemo_entity, pywemo_
     assert hass.states.get(wemo_entity.entity_id).state == STATE_UNAVAILABLE
 
     # Check that the entity recovers and is available after the update succeeds.
-    await hass.services.async_call(
-        HA_DOMAIN,
-        SERVICE_UPDATE_ENTITY,
-        {ATTR_ENTITY_ID: [wemo_entity.entity_id]},
-        blocking=True,
-    )
-
+    event.set()
+    await hass.async_block_till_done()
     assert hass.states.get(wemo_entity.entity_id).state == STATE_OFF
diff --git a/tests/components/wemo/test_light_bridge.py b/tests/components/wemo/test_light_bridge.py
index 3e7f79200c6..573f75a66d9 100644
--- a/tests/components/wemo/test_light_bridge.py
+++ b/tests/components/wemo/test_light_bridge.py
@@ -69,9 +69,10 @@ async def test_async_update_with_timeout_and_recovery(
     hass, pywemo_bridge_light, wemo_entity, pywemo_device
 ):
     """Test that the entity becomes unavailable after a timeout, and that it recovers."""
-    await entity_test_helpers.test_async_update_with_timeout_and_recovery(
-        hass, wemo_entity, pywemo_device
-    )
+    with _bypass_throttling():
+        await entity_test_helpers.test_async_update_with_timeout_and_recovery(
+            hass, wemo_entity, pywemo_device
+        )
 
 
 async def test_async_locked_update_with_exception(