Await callbacks to keep cleaner stacktraces (home-assistant#43693)

homeassistant/components/mqtt/__init__.py

@@ -1485,20 +1485,22 @@ def async_subscribe_connection_status(hass, connection_status_callback):
 
     connection_status_callback_job = HassJob(connection_status_callback)
 
-    @callback
-    def connected():
-        hass.async_add_hass_job(connection_status_callback_job, True)
+    async def connected():
+        task = hass.async_run_hass_job(connection_status_callback_job, True)
+        if task:
+            await task
 
-    @callback
-    def disconnected():
-        _LOGGER.error("Calling connection_status_callback, False")
-        hass.async_add_hass_job(connection_status_callback_job, False)
+    async def disconnected():
+        task = hass.async_run_hass_job(connection_status_callback_job, False)
+        if task:
+            await task
 
     subscriptions = {
         "connect": async_dispatcher_connect(hass, MQTT_CONNECTED, connected),
         "disconnect": async_dispatcher_connect(hass, MQTT_DISCONNECTED, disconnected),
     }
 
+    @callback
     def unsubscribe():
         subscriptions["connect"]()
         subscriptions["disconnect"]()

homeassistant/helpers/debounce.py

@@ -48,7 +48,7 @@ def function(self, function: Callable[..., Awaitable[Any]]) -> None:
 
     async def async_call(self) -> None:
         """Call the function."""
-        assert self.function is not None
+        assert self._job is not None
 
         if self._timer_task:
             if not self._execute_at_end_of_timer:
@@ -70,13 +70,15 @@ async def async_call(self) -> None:
             if self._timer_task:
                 return
 
-            await self.hass.async_add_hass_job(self._job)  # type: ignore
+            task = self.hass.async_run_hass_job(self._job)
+            if task:
+                await task
 
             self._schedule_timer()
 
     async def _handle_timer_finish(self) -> None:
         """Handle a finished timer."""
-        assert self.function is not None
+        assert self._job is not None
 
         self._timer_task = None
 
@@ -95,7 +97,9 @@ async def _handle_timer_finish(self) -> None:
                 return  # type: ignore
 
             try:
-                await self.hass.async_add_hass_job(self._job)  # type: ignore
+                task = self.hass.async_run_hass_job(self._job)
+                if task:
+                    await task
             except Exception:  # pylint: disable=broad-except
                 self.logger.exception("Unexpected exception from %s", self.function)
 

homeassistant/helpers/discovery.py

@@ -44,13 +44,14 @@ def async_listen(
 
     job = core.HassJob(callback)
 
-    @core.callback
-    def discovery_event_listener(event: core.Event) -> None:
+    async def discovery_event_listener(event: core.Event) -> None:
         """Listen for discovery events."""
         if ATTR_SERVICE in event.data and event.data[ATTR_SERVICE] in service:
-            hass.async_add_hass_job(
+            task = hass.async_run_hass_job(
                 job, event.data[ATTR_SERVICE], event.data.get(ATTR_DISCOVERED)
             )
+            if task:
+                await task
 
     hass.bus.async_listen(EVENT_PLATFORM_DISCOVERED, discovery_event_listener)
 
@@ -114,8 +115,7 @@ def async_listen_platform(
     service = EVENT_LOAD_PLATFORM.format(component)
     job = core.HassJob(callback)
 
-    @core.callback
-    def discovery_platform_listener(event: core.Event) -> None:
+    async def discovery_platform_listener(event: core.Event) -> None:
         """Listen for platform discovery events."""
         if event.data.get(ATTR_SERVICE) != service:
             return
@@ -125,7 +125,9 @@ def discovery_platform_listener(event: core.Event) -> None:
         if not platform:
             return
 
-        hass.async_run_hass_job(job, platform, event.data.get(ATTR_DISCOVERED))
+        task = hass.async_run_hass_job(job, platform, event.data.get(ATTR_DISCOVERED))
+        if task:
+            await task
 
     hass.bus.async_listen(EVENT_PLATFORM_DISCOVERED, discovery_platform_listener)
 

tests/test_loader.py

@@ -1,9 +1,9 @@
 """Test to verify that we can load components."""
 import pytest
 
+from homeassistant import core, loader
 from homeassistant.components import http, hue
 from homeassistant.components.hue import light as hue_light
-import homeassistant.loader as loader
 
 from tests.async_mock import ANY, patch
 from tests.common import MockModule, async_mock_service, mock_integration
@@ -83,6 +83,7 @@ async def test_helpers_wrapper(hass):
 
     result = []
 
+    @core.callback
     def discovery_callback(service, discovered):
         """Handle discovery callback."""
         result.append(discovered)