Switch to asyncio.wait for slow update warning implementation (#41184)

homeassistant/components/seventeentrack/sensor.py

@@ -16,6 +16,7 @@ from homeassistant.const import (
 )
 from homeassistant.helpers import aiohttp_client, config_validation as cv
 from homeassistant.helpers.entity import Entity
+from homeassistant.helpers.event import async_call_later
 from homeassistant.util import Throttle, slugify
 
 _LOGGER = logging.getLogger(__name__)
@@ -220,7 +221,8 @@ class SeventeenTrackPackageSensor(Entity):
         await self._data.async_update()
 
         if not self.available:
-            self.hass.async_create_task(self._remove())
+            # Entity cannot be removed while its being added
+            async_call_later(self.hass, 1, self._remove)
             return
 
         package = self._data.packages.get(self._tracking_number, None)
@@ -229,7 +231,8 @@ class SeventeenTrackPackageSensor(Entity):
         # delivered, post a notification:
         if package.status == VALUE_DELIVERED and not self._data.show_delivered:
             self._notify_delivered()
-            self.hass.async_create_task(self._remove())
+            # Entity cannot be removed while its being added
+            async_call_later(self.hass, 1, self._remove)
             return
 
         self._attrs.update(
@@ -238,7 +241,7 @@ class SeventeenTrackPackageSensor(Entity):
         self._state = package.status
         self._friendly_name = package.friendly_name
 
-    async def _remove(self):
+    async def _remove(self, *_):
         """Remove entity itself."""
         await self.async_remove()
 

homeassistant/helpers/entity.py

@@ -453,26 +453,35 @@ class Entity(ABC):
         if self.parallel_updates:
             await self.parallel_updates.acquire()
 
-        assert self.hass is not None
+        try:
-        if warning:
+            # pylint: disable=no-member
-            update_warn = self.hass.loop.call_later(
+            if hasattr(self, "async_update"):
-                SLOW_UPDATE_WARNING,
+                task = self.hass.async_create_task(self.async_update())  # type: ignore
-                _LOGGER.warning,
+            elif hasattr(self, "update"):
+                task = self.hass.async_add_executor_job(self.update)  # type: ignore
+            else:
+                return
+
+            if not warning:
+                await task
+                return
+
+            finished, _ = await asyncio.wait([task], timeout=SLOW_UPDATE_WARNING)
+
+            for done in finished:
+                exc = done.exception()
+                if exc:
+                    raise exc
+                return
+
+            _LOGGER.warning(
                 "Update of %s is taking over %s seconds",
                 self.entity_id,
                 SLOW_UPDATE_WARNING,
             )
-
+            await task
-        try:
-            # pylint: disable=no-member
-            if hasattr(self, "async_update"):
-                await self.async_update()  # type: ignore
-            elif hasattr(self, "update"):
-                await self.hass.async_add_executor_job(self.update)  # type: ignore
         finally:
             self._update_staged = False
-            if warning:
-                update_warn.cancel()
             if self.parallel_updates:
                 self.parallel_updates.release()
 

tests/components/group/test_cover.py

@@ -494,6 +494,7 @@ async def test_is_opening_closing(hass, setup_comp):
     await hass.services.async_call(
         DOMAIN, SERVICE_OPEN_COVER, {ATTR_ENTITY_ID: COVER_GROUP}, blocking=True
     )
+    await hass.async_block_till_done()
 
     assert hass.states.get(DEMO_COVER_POS).state == STATE_OPENING
     assert hass.states.get(DEMO_COVER_TILT).state == STATE_OPENING

tests/components/heos/test_media_player.py

@@ -120,6 +120,7 @@ async def test_updates_from_signals(hass, config_entry, config, controller, favo
         const.SIGNAL_PLAYER_EVENT, player.player_id, const.EVENT_PLAYER_STATE_CHANGED
     )
     await hass.async_block_till_done()
+
     state = hass.states.get("media_player.test_player")
     assert state.state == STATE_PLAYING
 
@@ -227,6 +228,7 @@ async def test_updates_from_players_changed(
         const.SIGNAL_CONTROLLER_EVENT, const.EVENT_PLAYERS_CHANGED, change_data
     )
     await event.wait()
+    await hass.async_block_till_done()
     assert hass.states.get("media_player.test_player").state == STATE_PLAYING
 
 

tests/components/kodi/test_device_trigger.py

@@ -112,6 +112,7 @@ async def test_if_fires_on_state_change(hass, calls, kodi_media_player):
             ]
         },
     )
+    await hass.async_block_till_done()
 
     await hass.services.async_call(
         MP_DOMAIN,

tests/components/seventeentrack/test_sensor.py

@@ -140,6 +140,8 @@ async def _goto_future(hass, future=None):
     with patch("homeassistant.util.utcnow", return_value=future):
         async_fire_time_changed(hass, future)
         await hass.async_block_till_done()
+        async_fire_time_changed(hass, future)
+        await hass.async_block_till_done()
 
 
 async def test_full_valid_config(hass):
@@ -247,6 +249,8 @@ async def test_delivered_not_shown(hass):
 
     hass.components.persistent_notification = MagicMock()
     await _setup_seventeentrack(hass, VALID_CONFIG_FULL_NO_DELIVERED)
+    await _goto_future(hass)
+
     assert not hass.states.async_entity_ids()
     hass.components.persistent_notification.create.assert_called()
 

tests/components/smart_meter_texas/test_init.py

@@ -45,6 +45,7 @@ async def test_update_failure(hass, config_entry, aioclient_mock):
     """Test that the coordinator handles a bad response."""
     await setup_integration(hass, config_entry, aioclient_mock, bad_reading=True)
     await async_setup_component(hass, HA_DOMAIN, {})
+    await hass.async_block_till_done()
     with patch("smart_meter_texas.Meter.read_meter") as updater:
         await hass.services.async_call(
             HA_DOMAIN,

tests/components/switch/test_light.py

@@ -56,6 +56,7 @@ async def test_light_service_calls(hass):
     assert hass.states.get("light.light_switch").state == "on"
 
     await common.async_turn_off(hass, "light.light_switch")
+    await hass.async_block_till_done()
 
     assert hass.states.get("switch.decorative_lights").state == "off"
     assert hass.states.get("light.light_switch").state == "off"
@@ -74,11 +75,13 @@ async def test_switch_service_calls(hass):
     assert hass.states.get("light.light_switch").state == "on"
 
     await switch_common.async_turn_off(hass, "switch.decorative_lights")
+    await hass.async_block_till_done()
 
     assert hass.states.get("switch.decorative_lights").state == "off"
     assert hass.states.get("light.light_switch").state == "off"
 
     await switch_common.async_turn_on(hass, "switch.decorative_lights")
+    await hass.async_block_till_done()
 
     assert hass.states.get("switch.decorative_lights").state == "on"
     assert hass.states.get("light.light_switch").state == "on"

tests/helpers/test_entity.py

@@ -114,13 +114,14 @@ class TestHelpersEntity:
         assert state.attributes.get(ATTR_DEVICE_CLASS) == "test_class"
 
 
-async def test_warn_slow_update(hass):
+async def test_warn_slow_update(hass, caplog):
     """Warn we log when entity update takes a long time."""
     update_call = False
 
     async def async_update():
         """Mock async update."""
         nonlocal update_call
+        await asyncio.sleep(0.00001)
         update_call = True
 
     mock_entity = entity.Entity()
@@ -128,22 +129,16 @@ async def test_warn_slow_update(hass):
     mock_entity.entity_id = "comp_test.test_entity"
     mock_entity.async_update = async_update
 
-    with patch.object(hass.loop, "call_later") as mock_call:
+    fast_update_time = 0.0000001
+
+    with patch.object(entity, "SLOW_UPDATE_WARNING", fast_update_time):
         await mock_entity.async_update_ha_state(True)
-        assert mock_call.called
+        assert str(fast_update_time) in caplog.text
-        assert len(mock_call.mock_calls) == 2
+        assert mock_entity.entity_id in caplog.text
-
-        timeout, logger_method = mock_call.mock_calls[0][1][:2]
-
-        assert timeout == entity.SLOW_UPDATE_WARNING
-        assert logger_method == entity._LOGGER.warning
-
-        assert mock_call().cancel.called
-
         assert update_call
 
 
-async def test_warn_slow_update_with_exception(hass):
+async def test_warn_slow_update_with_exception(hass, caplog):
     """Warn we log when entity update takes a long time and trow exception."""
     update_call = False
 
@@ -151,6 +146,7 @@ async def test_warn_slow_update_with_exception(hass):
         """Mock async update."""
         nonlocal update_call
         update_call = True
+        await asyncio.sleep(0.00001)
         raise AssertionError("Fake update error")
 
     mock_entity = entity.Entity()
@@ -158,28 +154,23 @@ async def test_warn_slow_update_with_exception(hass):
     mock_entity.entity_id = "comp_test.test_entity"
     mock_entity.async_update = async_update
 
-    with patch.object(hass.loop, "call_later") as mock_call:
+    fast_update_time = 0.0000001
+
+    with patch.object(entity, "SLOW_UPDATE_WARNING", fast_update_time):
         await mock_entity.async_update_ha_state(True)
-        assert mock_call.called
+        assert str(fast_update_time) in caplog.text
-        assert len(mock_call.mock_calls) == 2
+        assert mock_entity.entity_id in caplog.text
-
-        timeout, logger_method = mock_call.mock_calls[0][1][:2]
-
-        assert timeout == entity.SLOW_UPDATE_WARNING
-        assert logger_method == entity._LOGGER.warning
-
-        assert mock_call().cancel.called
-
         assert update_call
 
 
-async def test_warn_slow_device_update_disabled(hass):
+async def test_warn_slow_device_update_disabled(hass, caplog):
     """Disable slow update warning with async_device_update."""
     update_call = False
 
     async def async_update():
         """Mock async update."""
         nonlocal update_call
+        await asyncio.sleep(0.00001)
         update_call = True
 
     mock_entity = entity.Entity()
@@ -187,10 +178,12 @@ async def test_warn_slow_device_update_disabled(hass):
     mock_entity.entity_id = "comp_test.test_entity"
     mock_entity.async_update = async_update
 
-    with patch.object(hass.loop, "call_later") as mock_call:
+    fast_update_time = 0.0000001
-        await mock_entity.async_device_update(warning=False)
 
-        assert not mock_call.called
+    with patch.object(entity, "SLOW_UPDATE_WARNING", fast_update_time):
+        await mock_entity.async_device_update(warning=False)
+        assert str(fast_update_time) not in caplog.text
+        assert mock_entity.entity_id not in caplog.text
         assert update_call