From b0adaece257bba74125dc5d292c90c3307897f76 Mon Sep 17 00:00:00 2001
From: puddly <32534428+puddly@users.noreply.github.com>
Date: Sat, 13 Jan 2024 16:32:25 -0500
Subject: [PATCH] Reload ZHA only a single time when the connection is lost
 multiple times (#107963)

* Reload only a single time when the connection is lost multiple times

* Ignore when reset task finishes, allow only one reset per `ZHAGateway`
---
 homeassistant/components/zha/core/gateway.py | 15 ++++++++--
 tests/components/zha/test_gateway.py         | 30 ++++++++++++++++++++
 2 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/homeassistant/components/zha/core/gateway.py b/homeassistant/components/zha/core/gateway.py
index 12e439f1059..3efdc77934a 100644
--- a/homeassistant/components/zha/core/gateway.py
+++ b/homeassistant/components/zha/core/gateway.py
@@ -142,7 +142,9 @@ class ZHAGateway:
         self._log_relay_handler = LogRelayHandler(hass, self)
         self.config_entry = config_entry
         self._unsubs: list[Callable[[], None]] = []
+
         self.shutting_down = False
+        self._reload_task: asyncio.Task | None = None
 
     def get_application_controller_data(self) -> tuple[ControllerApplication, dict]:
         """Get an uninitialized instance of a zigpy `ControllerApplication`."""
@@ -231,12 +233,17 @@ class ZHAGateway:
 
     def connection_lost(self, exc: Exception) -> None:
         """Handle connection lost event."""
+        _LOGGER.debug("Connection to the radio was lost: %r", exc)
+
         if self.shutting_down:
             return
 
-        _LOGGER.debug("Connection to the radio was lost: %r", exc)
+        # Ensure we do not queue up multiple resets
+        if self._reload_task is not None:
+            _LOGGER.debug("Ignoring reset, one is already running")
+            return
 
-        self.hass.async_create_task(
+        self._reload_task = self.hass.async_create_task(
             self.hass.config_entries.async_reload(self.config_entry.entry_id)
         )
 
@@ -760,6 +767,10 @@ class ZHAGateway:
 
     async def shutdown(self) -> None:
         """Stop ZHA Controller Application."""
+        if self.shutting_down:
+            _LOGGER.debug("Ignoring duplicate shutdown event")
+            return
+
         _LOGGER.debug("Shutting down ZHA ControllerApplication")
         self.shutting_down = True
 
diff --git a/tests/components/zha/test_gateway.py b/tests/components/zha/test_gateway.py
index 4f520920704..9c3cf7aa2f8 100644
--- a/tests/components/zha/test_gateway.py
+++ b/tests/components/zha/test_gateway.py
@@ -291,3 +291,33 @@ async def test_gateway_force_multi_pan_channel(
 
     _, config = zha_gateway.get_application_controller_data()
     assert config["network"]["channel"] == expected_channel
+
+
+async def test_single_reload_on_multiple_connection_loss(
+    hass: HomeAssistant,
+    zigpy_app_controller: ControllerApplication,
+    config_entry: MockConfigEntry,
+):
+    """Test that we only reload once when we lose the connection multiple times."""
+    config_entry.add_to_hass(hass)
+
+    zha_gateway = ZHAGateway(hass, {}, config_entry)
+
+    with patch(
+        "bellows.zigbee.application.ControllerApplication.new",
+        return_value=zigpy_app_controller,
+    ):
+        await zha_gateway.async_initialize()
+
+        with patch.object(
+            hass.config_entries, "async_reload", wraps=hass.config_entries.async_reload
+        ) as mock_reload:
+            zha_gateway.connection_lost(RuntimeError())
+            zha_gateway.connection_lost(RuntimeError())
+            zha_gateway.connection_lost(RuntimeError())
+            zha_gateway.connection_lost(RuntimeError())
+            zha_gateway.connection_lost(RuntimeError())
+
+        assert len(mock_reload.mock_calls) == 1
+
+        await hass.async_block_till_done()