Add async_schedule_reload helper to the ConfigEntries manager (#110912)

* Add async_schedule_reload helper to the ConfigEntries manager

We have cases where the the setup retry kicks in right before
the reload happens causing the reload to fail with
OperationNotAllowed. The async_schedule_reload will
cancel the setup retry before the async_reload task
is created to avoid this problem.

I updated a few integrations that were most likely
to have this problem. Future PRs will do a more
extensive audit

* coverage

* revert for now since this needs more refactoring in a followup

* cover

* cleanup and fixes

homeassistant/components/flux_led/config_flow.py

@@ -125,9 +125,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
                     config_entries.ConfigEntryState.NOT_LOADED,
                 )
             ) or entry.state == config_entries.ConfigEntryState.SETUP_RETRY:
-                self.hass.async_create_task(
-                    self.hass.config_entries.async_reload(entry.entry_id)
-                )
+                self.hass.config_entries.async_schedule_reload(entry.entry_id)
             else:
                 async_dispatcher_send(
                     self.hass,

homeassistant/components/powerwall/config_flow.py

@@ -125,18 +125,14 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
                     if self.hass.config_entries.async_update_entry(
                         entry, unique_id=gateway_din
                     ):
-                        self.hass.async_create_task(
-                            self.hass.config_entries.async_reload(entry.entry_id)
-                        )
+                        self.hass.config_entries.async_schedule_reload(entry.entry_id)
                 return self.async_abort(reason="already_configured")
             if entry.unique_id == gateway_din:
                 if await self._async_powerwall_is_offline(entry):
                     if self.hass.config_entries.async_update_entry(
                         entry, data={**entry.data, CONF_IP_ADDRESS: self.ip_address}
                     ):
-                        self.hass.async_create_task(
-                            self.hass.config_entries.async_reload(entry.entry_id)
-                        )
+                        self.hass.config_entries.async_schedule_reload(entry.entry_id)
                 return self.async_abort(reason="already_configured")
         # Still need to abort for ignored entries
         self._abort_if_unique_id_configured()

homeassistant/components/tplink/config_flow.py

@@ -28,7 +28,7 @@ from homeassistant.const import (
     CONF_USERNAME,
 )
 from homeassistant.core import callback
-from homeassistant.data_entry_flow import AbortFlow, FlowResult
+from homeassistant.data_entry_flow import FlowResult
 from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.typing import DiscoveryInfoType
 
@@ -77,25 +77,22 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
     @callback
     def _update_config_if_entry_in_setup_error(
         self, entry: ConfigEntry, host: str, config: dict
-    ) -> None:
+    ) -> FlowResult | None:
         """If discovery encounters a device that is in SETUP_ERROR or SETUP_RETRY update the device config."""
         if entry.state not in (
             ConfigEntryState.SETUP_ERROR,
             ConfigEntryState.SETUP_RETRY,
         ):
-            return
+            return None
         entry_data = entry.data
         entry_config_dict = entry_data.get(CONF_DEVICE_CONFIG)
         if entry_config_dict == config and entry_data[CONF_HOST] == host:
-            return
-        self.hass.config_entries.async_update_entry(
-            entry, data={**entry.data, CONF_DEVICE_CONFIG: config, CONF_HOST: host}
+            return None
+        return self.async_update_reload_and_abort(
+            entry,
+            data={**entry.data, CONF_DEVICE_CONFIG: config, CONF_HOST: host},
+            reason="already_configured",
         )
-        self.hass.async_create_task(
-            self.hass.config_entries.async_reload(entry.entry_id),
-            f"config entry reload {entry.title} {entry.domain} {entry.entry_id}",
-        )
-        raise AbortFlow("already_configured")
 
     async def _async_handle_discovery(
         self, host: str, formatted_mac: str, config: dict | None = None
@@ -104,8 +101,16 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         current_entry = await self.async_set_unique_id(
             formatted_mac, raise_on_progress=False
         )
-        if config and current_entry:
-            self._update_config_if_entry_in_setup_error(current_entry, host, config)
+        if (
+            config
+            and current_entry
+            and (
+                result := self._update_config_if_entry_in_setup_error(
+                    current_entry, host, config
+                )
+            )
+        ):
+            return result
         self._abort_if_unique_id_configured(updates={CONF_HOST: host})
         self._async_abort_entries_match({CONF_HOST: host})
         self.context[CONF_HOST] = host

homeassistant/components/upnp/config_flow.py

@@ -206,20 +206,12 @@ class UpnpFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
                 return self.async_abort(reason="discovery_ignored")
 
             LOGGER.debug("Updating entry: %s", entry.entry_id)
-            self.hass.config_entries.async_update_entry(
+            return self.async_update_reload_and_abort(
                 entry,
                 unique_id=unique_id,
                 data={**entry.data, CONFIG_ENTRY_UDN: discovery_info.ssdp_udn},
+                reason="config_entry_updated",
             )
-            if entry.state == config_entries.ConfigEntryState.LOADED:
-                # Only reload when entry has state LOADED; when entry has state
-                # SETUP_RETRY, another load is started,
-                # causing the entry to be loaded twice.
-                LOGGER.debug("Reloading entry: %s", entry.entry_id)
-                self.hass.async_create_task(
-                    self.hass.config_entries.async_reload(entry.entry_id)
-                )
-            return self.async_abort(reason="config_entry_updated")
 
         # Store discovery.
         self._add_discovery(discovery_info)

homeassistant/components/yeelight/config_flow.py

@@ -102,9 +102,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
                     ConfigEntryState.LOADED,
                 )
             if reload:
-                self.hass.async_create_task(
-                    self.hass.config_entries.async_reload(entry.entry_id)
-                )
+                self.hass.config_entries.async_schedule_reload(entry.entry_id)
             return self.async_abort(reason="already_configured")
         return await self._async_handle_discovery()
 

homeassistant/components/zwave_js/config_flow.py

@@ -750,9 +750,7 @@ class OptionsFlowHandler(BaseZwaveJSFlow, config_entries.OptionsFlow):
                 }
             )
 
-            self.hass.async_create_task(
-                self.hass.config_entries.async_reload(self.config_entry.entry_id)
-            )
+            self.hass.config_entries.async_schedule_reload(self.config_entry.entry_id)
             return self.async_create_entry(title=TITLE, data={})
 
         return self.async_show_form(
@@ -917,9 +915,7 @@ class OptionsFlowHandler(BaseZwaveJSFlow, config_entries.OptionsFlow):
             }
         )
         # Always reload entry since we may have disconnected the client.
-        self.hass.async_create_task(
-            self.hass.config_entries.async_reload(self.config_entry.entry_id)
-        )
+        self.hass.config_entries.async_schedule_reload(self.config_entry.entry_id)
         return self.async_create_entry(title=TITLE, data={})
 
     async def async_revert_addon_config(self, reason: str) -> FlowResult:
@@ -935,9 +931,7 @@ class OptionsFlowHandler(BaseZwaveJSFlow, config_entries.OptionsFlow):
             )
 
         if self.revert_reason or not self.original_addon_config:
-            self.hass.async_create_task(
-                self.hass.config_entries.async_reload(self.config_entry.entry_id)
-            )
+            self.hass.config_entries.async_schedule_reload(self.config_entry.entry_id)
             return self.async_abort(reason=reason)
 
         self.revert_reason = reason

homeassistant/config_entries.py

@@ -1519,6 +1519,17 @@ class ConfigEntries:
 
         return await entry.async_unload(self.hass)
 
+    @callback
+    def async_schedule_reload(self, entry_id: str) -> None:
+        """Schedule a config entry to be reloaded."""
+        if (entry := self.async_get_entry(entry_id)) is None:
+            raise UnknownEntry
+        entry.async_cancel_retry_setup()
+        self.hass.async_create_task(
+            self.async_reload(entry_id),
+            f"config entry reload {entry.title} {entry.domain} {entry.entry_id}",
+        )
+
     async def async_reload(self, entry_id: str) -> bool:
         """Reload an entry.
 
@@ -1861,10 +1872,7 @@ class ConfigFlow(data_entry_flow.FlowHandler):
         if entry.source == SOURCE_IGNORE and self.source == SOURCE_USER:
             return
         if should_reload:
-            self.hass.async_create_task(
-                self.hass.config_entries.async_reload(entry.entry_id),
-                f"config entry reload {entry.title} {entry.domain} {entry.entry_id}",
-            )
+            self.hass.config_entries.async_schedule_reload(entry.entry_id)
         raise data_entry_flow.AbortFlow(error)
 
     async def async_set_unique_id(
@@ -2123,10 +2131,7 @@ class ConfigFlow(data_entry_flow.FlowHandler):
             options=options,
         )
         if result:
-            self.hass.async_create_task(
-                self.hass.config_entries.async_reload(entry.entry_id),
-                f"config entry reload {entry.title} {entry.domain} {entry.entry_id}",
-            )
+            self.hass.config_entries.async_schedule_reload(entry.entry_id)
         return self.async_abort(reason=reason)
 
 

tests/test_config_entries.py

@@ -3602,6 +3602,39 @@ async def test_setup_retrying_during_shutdown(hass: HomeAssistant) -> None:
     entry.async_cancel_retry_setup()
 
 
+async def test_scheduling_reload_cancels_setup_retry(hass: HomeAssistant) -> None:
+    """Test scheduling a reload cancels setup retry."""
+    entry = MockConfigEntry(domain="test")
+    entry.add_to_hass(hass)
+
+    mock_setup_entry = AsyncMock(side_effect=ConfigEntryNotReady)
+    mock_integration(hass, MockModule("test", async_setup_entry=mock_setup_entry))
+    mock_platform(hass, "test.config_flow", None)
+    cancel_mock = Mock()
+
+    with patch(
+        "homeassistant.config_entries.async_call_later", return_value=cancel_mock
+    ):
+        await entry.async_setup(hass)
+
+    assert entry.state is config_entries.ConfigEntryState.SETUP_RETRY
+    assert len(cancel_mock.mock_calls) == 0
+
+    mock_setup_entry.side_effect = None
+    mock_setup_entry.return_value = True
+    hass.config_entries.async_schedule_reload(entry.entry_id)
+
+    assert len(cancel_mock.mock_calls) == 1
+    await hass.async_block_till_done()
+    assert entry.state is config_entries.ConfigEntryState.LOADED
+
+
+async def test_scheduling_reload_unknown_entry(hass: HomeAssistant) -> None:
+    """Test scheduling a reload raises with an unknown entry."""
+    with pytest.raises(config_entries.UnknownEntry):
+        hass.config_entries.async_schedule_reload("non-existing")
+
+
 @pytest.mark.parametrize(
     ("matchers", "reason"),
     [
@@ -3947,7 +3980,7 @@ async def test_unique_id_update_while_setup_in_progress(
     assert entry.state is config_entries.ConfigEntryState.LOADED
 
 
-async def test_disallow_entry_reload_with_setup_in_progresss(
+async def test_disallow_entry_reload_with_setup_in_progress(
     hass: HomeAssistant, manager: config_entries.ConfigEntries
 ) -> None:
     """Test we do not allow reload while the config entry is still setting up."""