diff --git a/homeassistant/config_entries.py b/homeassistant/config_entries.py index 75b0631339f..e48313cab33 100644 --- a/homeassistant/config_entries.py +++ b/homeassistant/config_entries.py @@ -3,7 +3,7 @@ from __future__ import annotations import asyncio -from collections import UserDict +from collections import UserDict, defaultdict from collections.abc import ( Callable, Coroutine, @@ -1224,8 +1224,12 @@ class ConfigEntriesFlowManager(data_entry_flow.FlowManager[ConfigFlowResult]): super().__init__(hass) self.config_entries = config_entries self._hass_config = hass_config - self._pending_import_flows: dict[str, dict[str, asyncio.Future[None]]] = {} - self._initialize_futures: dict[str, list[asyncio.Future[None]]] = {} + self._pending_import_flows: defaultdict[ + str, dict[str, asyncio.Future[None]] + ] = defaultdict(dict) + self._initialize_futures: defaultdict[str, set[asyncio.Future[None]]] = ( + defaultdict(set) + ) self._discovery_debouncer = Debouncer[None]( hass, _LOGGER, @@ -1278,12 +1282,11 @@ class ConfigEntriesFlowManager(data_entry_flow.FlowManager[ConfigFlowResult]): loop = self.hass.loop if context["source"] == SOURCE_IMPORT: - self._pending_import_flows.setdefault(handler, {})[flow_id] = ( - loop.create_future() - ) + self._pending_import_flows[handler][flow_id] = loop.create_future() cancel_init_future = loop.create_future() - self._initialize_futures.setdefault(handler, []).append(cancel_init_future) + handler_init_futures = self._initialize_futures[handler] + handler_init_futures.add(cancel_init_future) try: async with interrupt( cancel_init_future, @@ -1294,8 +1297,13 @@ class ConfigEntriesFlowManager(data_entry_flow.FlowManager[ConfigFlowResult]): except FlowCancelledError as ex: raise asyncio.CancelledError from ex finally: - self._initialize_futures[handler].remove(cancel_init_future) - self._pending_import_flows.get(handler, {}).pop(flow_id, None) + handler_init_futures.remove(cancel_init_future) + if not handler_init_futures: + del self._initialize_futures[handler] + if handler in self._pending_import_flows: + self._pending_import_flows[handler].pop(flow_id, None) + if not self._pending_import_flows[handler]: + del self._pending_import_flows[handler] if result["type"] != data_entry_flow.FlowResultType.ABORT: await self.async_post_init(flow, result) @@ -1322,11 +1330,18 @@ class ConfigEntriesFlowManager(data_entry_flow.FlowManager[ConfigFlowResult]): try: result = await self._async_handle_step(flow, flow.init_step, data) finally: - init_done = self._pending_import_flows.get(handler, {}).get(flow_id) - if init_done and not init_done.done(): - init_done.set_result(None) + self._set_pending_import_done(flow) return flow, result + def _set_pending_import_done(self, flow: ConfigFlow) -> None: + """Set pending import flow as done.""" + if ( + (handler_import_flows := self._pending_import_flows.get(flow.handler)) + and (init_done := handler_import_flows.get(flow.flow_id)) + and not init_done.done() + ): + init_done.set_result(None) + @callback def async_shutdown(self) -> None: """Cancel any initializing flows.""" @@ -1347,9 +1362,7 @@ class ConfigEntriesFlowManager(data_entry_flow.FlowManager[ConfigFlowResult]): # We do this to avoid a circular dependency where async_finish_flow sets up a # new entry, which needs the integration to be set up, which is waiting for # init to be done. - init_done = self._pending_import_flows.get(flow.handler, {}).get(flow.flow_id) - if init_done and not init_done.done(): - init_done.set_result(None) + self._set_pending_import_done(flow) # Remove notification if no other discovery config entries in progress if not self._async_has_other_discovery_flows(flow.flow_id):