From 3390fb32a8a610318e55a832fab65a47cfd37a6b Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Thu, 6 Feb 2025 15:18:37 +0100
Subject: [PATCH] Don't overwrite setup state in async_set_domains_to_be_loaded
 (#137547)

---
 homeassistant/setup.py |  8 ++++-
 tests/test_setup.py    | 76 +++++++++++++++++++++++++++++++++++++++---
 2 files changed, 78 insertions(+), 6 deletions(-)

diff --git a/homeassistant/setup.py b/homeassistant/setup.py
index 1fa93a80cd5..dc4d0988b91 100644
--- a/homeassistant/setup.py
+++ b/homeassistant/setup.py
@@ -132,7 +132,13 @@ def async_set_domains_to_be_loaded(hass: core.HomeAssistant, domains: set[str])
      - Keep track of domains which will load but have not yet finished loading
     """
     setup_done_futures = hass.data.setdefault(DATA_SETUP_DONE, {})
-    setup_done_futures.update({domain: hass.loop.create_future() for domain in domains})
+    setup_futures = hass.data.setdefault(DATA_SETUP, {})
+    old_domains = set(setup_futures) | set(setup_done_futures) | hass.config.components
+    if overlap := old_domains & domains:
+        _LOGGER.debug("Domains to be loaded %s already loaded or pending", overlap)
+    setup_done_futures.update(
+        {domain: hass.loop.create_future() for domain in domains - old_domains}
+    )
 
 
 def setup_component(hass: core.HomeAssistant, domain: str, config: ConfigType) -> bool:
diff --git a/tests/test_setup.py b/tests/test_setup.py
index 2d15c670cf7..bb221c7cb4c 100644
--- a/tests/test_setup.py
+++ b/tests/test_setup.py
@@ -363,20 +363,24 @@ async def test_component_failing_setup(hass: HomeAssistant) -> None:
 
 async def test_component_exception_setup(hass: HomeAssistant) -> None:
     """Test component that raises exception during setup."""
-    setup.async_set_domains_to_be_loaded(hass, {"comp"})
+    domain = "comp"
+    setup.async_set_domains_to_be_loaded(hass, {domain})
 
     def exception_setup(hass: HomeAssistant, config: ConfigType) -> bool:
         """Raise exception."""
         raise Exception("fail!")  # noqa: TRY002
 
-    mock_integration(hass, MockModule("comp", setup=exception_setup))
+    mock_integration(hass, MockModule(domain, setup=exception_setup))
 
-    assert not await setup.async_setup_component(hass, "comp", {})
-    assert "comp" not in hass.config.components
+    assert not await setup.async_setup_component(hass, domain, {})
+    assert domain in hass.data[setup.DATA_SETUP]
+    assert domain not in hass.data[setup.DATA_SETUP_DONE]
+    assert domain not in hass.config.components
 
 
 async def test_component_base_exception_setup(hass: HomeAssistant) -> None:
     """Test component that raises exception during setup."""
+    domain = "comp"
     setup.async_set_domains_to_be_loaded(hass, {"comp"})
 
     def exception_setup(hass: HomeAssistant, config: ConfigType) -> bool:
@@ -389,7 +393,69 @@ async def test_component_base_exception_setup(hass: HomeAssistant) -> None:
         await setup.async_setup_component(hass, "comp", {})
     assert str(exc_info.value) == "fail!"
 
-    assert "comp" not in hass.config.components
+    assert domain in hass.data[setup.DATA_SETUP]
+    assert domain not in hass.data[setup.DATA_SETUP_DONE]
+    assert domain not in hass.config.components
+
+
+async def test_set_domains_to_be_loaded(hass: HomeAssistant) -> None:
+    """Test async_set_domains_to_be_loaded."""
+    domain_good = "comp_good"
+    domain_bad = "comp_bad"
+    domain_base_exception = "comp_base_exception"
+    domain_exception = "comp_exception"
+    domains = {domain_good, domain_bad, domain_exception, domain_base_exception}
+    setup.async_set_domains_to_be_loaded(hass, domains)
+
+    assert set(hass.data[setup.DATA_SETUP_DONE]) == domains
+    setup_done = dict(hass.data[setup.DATA_SETUP_DONE])
+
+    # Calling async_set_domains_to_be_loaded again should not create new futures
+    setup.async_set_domains_to_be_loaded(hass, domains)
+    assert setup_done == hass.data[setup.DATA_SETUP_DONE]
+
+    def good_setup(hass: HomeAssistant, config: ConfigType) -> bool:
+        """Success."""
+        return True
+
+    def bad_setup(hass: HomeAssistant, config: ConfigType) -> bool:
+        """Fail."""
+        return False
+
+    def base_exception_setup(hass: HomeAssistant, config: ConfigType) -> bool:
+        """Raise exception."""
+        raise BaseException("fail!")  # noqa: TRY002
+
+    def exception_setup(hass: HomeAssistant, config: ConfigType) -> bool:
+        """Raise exception."""
+        raise Exception("fail!")  # noqa: TRY002
+
+    mock_integration(hass, MockModule(domain_good, setup=good_setup))
+    mock_integration(hass, MockModule(domain_bad, setup=bad_setup))
+    mock_integration(
+        hass, MockModule(domain_base_exception, setup=base_exception_setup)
+    )
+    mock_integration(hass, MockModule(domain_exception, setup=exception_setup))
+
+    # Set up the four components
+    assert await setup.async_setup_component(hass, domain_good, {})
+    assert not await setup.async_setup_component(hass, domain_bad, {})
+    assert not await setup.async_setup_component(hass, domain_exception, {})
+    with pytest.raises(BaseException, match="fail!"):
+        await setup.async_setup_component(hass, domain_base_exception, {})
+
+    # Check the result of the setup
+    assert not hass.data[setup.DATA_SETUP_DONE]
+    assert set(hass.data[setup.DATA_SETUP]) == {
+        domain_bad,
+        domain_exception,
+        domain_base_exception,
+    }
+    assert set(hass.config.components) == {domain_good}
+
+    # Calling async_set_domains_to_be_loaded again should not create any new futures
+    setup.async_set_domains_to_be_loaded(hass, domains)
+    assert not hass.data[setup.DATA_SETUP_DONE]
 
 
 async def test_component_setup_with_validation_and_dependency(