Fix cancel propagation in update coordinator and config entry (home-assistant#153504)

Author: elupus

homeassistant/components/history_stats/config_flow.py

@@ -211,7 +211,7 @@ async def ws_start_preview(
 
     @callback
     def async_preview_updated(
-        last_exception: Exception | None, state: str, attributes: Mapping[str, Any]
+        last_exception: BaseException | None, state: str, attributes: Mapping[str, Any]
     ) -> None:
         """Forward config entry state events to websocket."""
         if last_exception:

homeassistant/components/history_stats/sensor.py

@@ -241,7 +241,9 @@ def _process_update(self) -> None:
 
     async def async_start_preview(
         self,
-        preview_callback: Callable[[Exception | None, str, Mapping[str, Any]], None],
+        preview_callback: Callable[
+            [BaseException | None, str, Mapping[str, Any]], None
+        ],
     ) -> CALLBACK_TYPE:
         """Render a preview."""
 

homeassistant/config_entries.py

@@ -754,6 +754,7 @@ async def __async_setup_with_context(
         error_reason_translation_key = None
         error_reason_translation_placeholders = None
 
+        result = False
         try:
             with async_start_setup(
                 hass, integration=self.domain, group=self.entry_id, phase=setup_phase
@@ -775,8 +776,6 @@ async def __async_setup_with_context(
                 self.domain,
                 error_reason,
             )
-            await self._async_process_on_unload(hass)
-            result = False
         except ConfigEntryAuthFailed as exc:
             message = str(exc)
             auth_base_message = "could not authenticate"
@@ -792,9 +791,7 @@ async def __async_setup_with_context(
                 self.domain,
                 auth_message,
             )
-            await self._async_process_on_unload(hass)
             self.async_start_reauth(hass)
-            result = False
         except ConfigEntryNotReady as exc:
             message = str(exc)
             error_reason_translation_key = exc.translation_key
@@ -835,14 +832,39 @@ async def __async_setup_with_context(
                     functools.partial(self._async_setup_again, hass),
                 )
 
-            await self._async_process_on_unload(hass)
             return
+
+        except asyncio.CancelledError:
+            # We want to propagate CancelledError if we are being cancelled.
+            if (task := asyncio.current_task()) and task.cancelling() > 0:
+                _LOGGER.exception(
+                    "Setup of config entry '%s' for %s integration cancelled",
+                    self.title,
+                    self.domain,
+                )
+                self._async_set_state(
+                    hass,
+                    ConfigEntryState.SETUP_ERROR,
+                    None,
+                    None,
+                    None,
+                )
+                raise
+
+            # This was not a "real" cancellation, log it and treat as a normal error.
+            _LOGGER.exception(
+                "Error setting up entry %s for %s", self.title, integration.domain
+            )
+
         # pylint: disable-next=broad-except
-        except (asyncio.CancelledError, SystemExit, Exception):
+        except (SystemExit, Exception):
             _LOGGER.exception(
                 "Error setting up entry %s for %s", self.title, integration.domain
             )
-            result = False
+
+        finally:
+            if not result and domain_is_integration:
+                await self._async_process_on_unload(hass)
 
         #
         # After successfully calling async_setup_entry, it is important that this function

homeassistant/helpers/update_coordinator.py

@@ -131,7 +131,7 @@ def __init__(
         self._request_refresh_task: asyncio.TimerHandle | None = None
         self._retry_after: float | None = None
         self.last_update_success = True
-        self.last_exception: Exception | None = None
+        self.last_exception: BaseException | None = None
 
         if request_refresh_debouncer is None:
             request_refresh_debouncer = Debouncer(
@@ -492,8 +492,16 @@ async def _async_refresh(  # noqa: C901
                 self.config_entry.async_start_reauth(self.hass)
         except NotImplementedError as err:
             self.last_exception = err
+            self.last_update_success = False
             raise
 
+        except asyncio.CancelledError as err:
+            self.last_exception = err
+            self.last_update_success = False
+
+            if (task := asyncio.current_task()) and task.cancelling() > 0:
+                raise
+
         except Exception as err:
             self.last_exception = err
             self.last_update_success = False

tests/helpers/test_update_coordinator.py

@@ -334,6 +334,35 @@ async def test_refresh_no_update_method(
         await crd.async_refresh()
 
 
+async def test_refresh_cancelled(
+    hass: HomeAssistant,
+    crd: update_coordinator.DataUpdateCoordinator[int],
+) -> None:
+    """Test that we don't swallow cancellation."""
+    await crd.async_refresh()
+
+    start = asyncio.Event()
+    abort = asyncio.Event()
+
+    async def _update() -> bool:
+        start.set()
+        await abort.wait()
+        return True
+
+    crd.update_method = _update
+    crd.last_update_success = True
+
+    task = hass.async_create_task(crd.async_refresh())
+    await start.wait()
+    task.cancel()
+
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    abort.set()
+    assert crd.last_update_success is False
+
+
 async def test_update_interval(
     hass: HomeAssistant,
     freezer: FrozenDateTimeFactory,

tests/test_config_entries.py

@@ -9727,3 +9727,31 @@ async def async_step_user(
     assert result["next_flow"][0] == config_entries.FlowType.CONFIG_FLOW
     # Verify the target flow exists
     hass.config_entries.flow.async_get(result["next_flow"][1])
+
+
+async def test_canceled_exceptions_are_propagated(
+    hass: HomeAssistant, manager: config_entries.ConfigEntries
+) -> None:
+    """Tests that base exceptions like cancellations are not swallowed."""
+    entry = MockConfigEntry(title="test_title", domain="test")
+    start = asyncio.Event()
+    abort = asyncio.Event()
+
+    async def _setup(_: HomeAssistant, __: ConfigEntry) -> bool:
+        start.set()
+        await abort.wait()
+        return True
+
+    mock_integration(hass, MockModule("test", async_setup_entry=_setup))
+    mock_platform(hass, "test.config_flow", None)
+
+    entry.add_to_hass(hass)
+
+    task = hass.async_create_task(manager.async_setup(entry.entry_id))
+    await start.wait()
+    task.cancel()
+
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    abort.set()

tests/testing_config/custom_components/test_package_raises_cancelled_error_config_entry/__init__.py

@@ -13,6 +13,5 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> None:
-    """Mock an unsuccessful entry setup."""
-    asyncio.current_task().cancel()
-    await asyncio.sleep(0)
+    """Mock an leaked cancellation, without our own task being cancelled."""
+    raise asyncio.CancelledError