From 2e47cee72a12253ee66d5d4bd1c1f2df47b000d5 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Mon, 6 Jun 2022 19:48:49 -1000
Subject: [PATCH] Fix setup race when config entry is in a setup retry state
 (#73145)

---
 homeassistant/config_entries.py |  6 +++
 tests/test_config_entries.py    | 72 ++++++++++++++++++++++++++++++++-
 2 files changed, 77 insertions(+), 1 deletion(-)

diff --git a/homeassistant/config_entries.py b/homeassistant/config_entries.py
index df633009138..bb24b24a7fb 100644
--- a/homeassistant/config_entries.py
+++ b/homeassistant/config_entries.py
@@ -90,6 +90,8 @@ class ConfigEntryState(Enum):
     """The config entry has not been loaded"""
     FAILED_UNLOAD = "failed_unload", False
     """An error occurred when trying to unload the entry"""
+    SETUP_IN_PROGRESS = "setup_in_progress", True
+    """The config entry is setting up."""
 
     _recoverable: bool
 
@@ -294,6 +296,10 @@ class ConfigEntry:
         if integration is None:
             integration = await loader.async_get_integration(hass, self.domain)
 
+        # Only store setup result as state if it was not forwarded.
+        if self.domain == integration.domain:
+            self.state = ConfigEntryState.SETUP_IN_PROGRESS
+
         self.supports_unload = await support_entry_unload(hass, self.domain)
         self.supports_remove_device = await support_remove_from_device(
             hass, self.domain
diff --git a/tests/test_config_entries.py b/tests/test_config_entries.py
index 09951b4f34e..dbbf542d36c 100644
--- a/tests/test_config_entries.py
+++ b/tests/test_config_entries.py
@@ -16,7 +16,7 @@ from homeassistant.const import (
     EVENT_HOMEASSISTANT_STARTED,
     EVENT_HOMEASSISTANT_STOP,
 )
-from homeassistant.core import CoreState, Event, callback
+from homeassistant.core import CoreState, Event, HomeAssistant, callback
 from homeassistant.data_entry_flow import RESULT_TYPE_ABORT, BaseServiceInfo, FlowResult
 from homeassistant.exceptions import (
     ConfigEntryAuthFailed,
@@ -3190,3 +3190,73 @@ async def test_entry_reload_concurrency(hass, manager):
     await asyncio.gather(*tasks)
     assert entry.state is config_entries.ConfigEntryState.LOADED
     assert loaded == 1
+
+
+async def test_unique_id_update_while_setup_in_progress(
+    hass: HomeAssistant, manager: config_entries.ConfigEntries
+) -> None:
+    """Test we handle the case where the config entry is updated while setup is in progress."""
+
+    async def mock_setup_entry(hass, entry):
+        """Mock setting up entry."""
+        await asyncio.sleep(0.1)
+        return True
+
+    async def mock_unload_entry(hass, entry):
+        """Mock unloading an entry."""
+        return True
+
+    hass.config.components.add("comp")
+    entry = MockConfigEntry(
+        domain="comp",
+        data={"additional": "data", "host": "0.0.0.0"},
+        unique_id="mock-unique-id",
+        state=config_entries.ConfigEntryState.SETUP_RETRY,
+    )
+    entry.add_to_hass(hass)
+
+    mock_integration(
+        hass,
+        MockModule(
+            "comp",
+            async_setup_entry=mock_setup_entry,
+            async_unload_entry=mock_unload_entry,
+        ),
+    )
+    mock_entity_platform(hass, "config_flow.comp", None)
+    updates = {"host": "1.1.1.1"}
+
+    hass.async_create_task(hass.config_entries.async_reload(entry.entry_id))
+    await asyncio.sleep(0)
+    assert entry.state is config_entries.ConfigEntryState.SETUP_IN_PROGRESS
+
+    class TestFlow(config_entries.ConfigFlow):
+        """Test flow."""
+
+        VERSION = 1
+
+        async def async_step_user(self, user_input=None):
+            """Test user step."""
+            await self.async_set_unique_id("mock-unique-id")
+            await self._abort_if_unique_id_configured(
+                updates=updates, reload_on_update=True
+            )
+
+    with patch.dict(config_entries.HANDLERS, {"comp": TestFlow}), patch(
+        "homeassistant.config_entries.ConfigEntries.async_reload"
+    ) as async_reload:
+        result = await manager.flow.async_init(
+            "comp", context={"source": config_entries.SOURCE_USER}
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == RESULT_TYPE_ABORT
+    assert result["reason"] == "already_configured"
+    assert entry.data["host"] == "1.1.1.1"
+    assert entry.data["additional"] == "data"
+
+    # Setup is already in progress, we should not reload
+    # if it fails it will go into a retry state and try again
+    assert len(async_reload.mock_calls) == 0
+    await hass.async_block_till_done()
+    assert entry.state is config_entries.ConfigEntryState.LOADED