Clean up stale ZHA database listener when reconnecting to radio (#101850)

homeassistant/components/zha/__init__.py

@@ -12,8 +12,8 @@ from zigpy.config import CONF_DATABASE, CONF_DEVICE, CONF_DEVICE_PATH
 from zigpy.exceptions import NetworkSettingsInconsistent
 
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_TYPE
+from homeassistant.const import CONF_TYPE, EVENT_HOMEASSISTANT_STOP
-from homeassistant.core import HomeAssistant
+from homeassistant.core import Event, HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady, HomeAssistantError
 from homeassistant.helpers import device_registry as dr
 import homeassistant.helpers.config_validation as cv
@@ -160,19 +160,6 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
 
     zha_gateway = ZHAGateway(hass, zha_data.yaml_config, config_entry)
 
-    async def async_zha_shutdown():
-        """Handle shutdown tasks."""
-        await zha_gateway.shutdown()
-        # clean up any remaining entity metadata
-        # (entities that have been discovered but not yet added to HA)
-        # suppress KeyError because we don't know what state we may
-        # be in when we get here in failure cases
-        with contextlib.suppress(KeyError):
-            for platform in PLATFORMS:
-                del zha_data.platforms[platform]
-
-    config_entry.async_on_unload(async_zha_shutdown)
-
     try:
         await zha_gateway.async_initialize()
     except NetworkSettingsInconsistent as exc:
@@ -211,6 +198,13 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
 
     websocket_api.async_load_api(hass)
 
+    async def async_shutdown(_: Event) -> None:
+        await zha_gateway.shutdown()
+
+    config_entry.async_on_unload(
+        hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, async_shutdown)
+    )
+
     await zha_gateway.async_initialize_devices_and_entities()
     await hass.config_entries.async_forward_entry_setups(config_entry, PLATFORMS)
     async_dispatcher_send(hass, SIGNAL_ADD_ENTITIES)
@@ -220,8 +214,19 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
 async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
     """Unload ZHA config entry."""
     zha_data = get_zha_data(hass)
+
+    if zha_data.gateway is not None:
+        await zha_data.gateway.shutdown()
         zha_data.gateway = None
 
+    # clean up any remaining entity metadata
+    # (entities that have been discovered but not yet added to HA)
+    # suppress KeyError because we don't know what state we may
+    # be in when we get here in failure cases
+    with contextlib.suppress(KeyError):
+        for platform in PLATFORMS:
+            del zha_data.platforms[platform]
+
     GROUP_PROBE.cleanup()
     websocket_api.async_unload_api(hass)
 

homeassistant/components/zha/core/gateway.py

@@ -203,15 +203,16 @@ class ZHAGateway:
             start_radio=False,
         )
 
+        try:
             for attempt in range(STARTUP_RETRIES):
                 try:
                     await self.application_controller.startup(auto_form=True)
-            except NetworkSettingsInconsistent:
-                raise
                 except TransientConnectionError as exc:
                     raise ConfigEntryNotReady from exc
+                except NetworkSettingsInconsistent:
+                    raise
                 except Exception as exc:  # pylint: disable=broad-except
-                _LOGGER.warning(
+                    _LOGGER.debug(
                         "Couldn't start %s coordinator (attempt %s of %s)",
                         self.radio_description,
                         attempt + 1,
@@ -225,6 +226,10 @@ class ZHAGateway:
                     await asyncio.sleep(STARTUP_FAILURE_DELAY_S)
                 else:
                     break
+        except Exception:
+            # Explicitly shut down the controller application on failure
+            await self.application_controller.shutdown()
+            raise
 
         zha_data = get_zha_data(self.hass)
         zha_data.gateway = self

homeassistant/components/zha/core/helpers.py

@@ -437,7 +437,10 @@ class ZHAData:
 
 def get_zha_data(hass: HomeAssistant) -> ZHAData:
     """Get the global ZHA data object."""
-    return hass.data.get(DATA_ZHA, ZHAData())
+    if DATA_ZHA not in hass.data:
+        hass.data[DATA_ZHA] = ZHAData()
+
+    return hass.data[DATA_ZHA]
 
 
 def get_zha_gateway(hass: HomeAssistant) -> ZHAGateway:

homeassistant/components/zha/manifest.json

@@ -26,7 +26,7 @@
     "pyserial-asyncio==0.6",
     "zha-quirks==0.0.105",
     "zigpy-deconz==0.21.1",
-    "zigpy==0.57.2",
+    "zigpy==0.58.1",
     "zigpy-xbee==0.18.3",
     "zigpy-zigate==0.11.0",
     "zigpy-znp==0.11.6",

requirements_all.txt

@@ -2813,7 +2813,7 @@ zigpy-zigate==0.11.0
 zigpy-znp==0.11.6
 
 # homeassistant.components.zha
-zigpy==0.57.2
+zigpy==0.58.1
 
 # homeassistant.components.zoneminder
 zm-py==0.5.2

requirements_test_all.txt

@@ -2101,7 +2101,7 @@ zigpy-zigate==0.11.0
 zigpy-znp==0.11.6
 
 # homeassistant.components.zha
-zigpy==0.57.2
+zigpy==0.58.1
 
 # homeassistant.components.zwave_js
 zwave-js-server-python==0.52.1

tests/components/zha/test_init.py

@@ -13,8 +13,14 @@ from homeassistant.components.zha.core.const import (
     CONF_USB_PATH,
     DOMAIN,
 )
-from homeassistant.const import MAJOR_VERSION, MINOR_VERSION, Platform
+from homeassistant.components.zha.core.helpers import get_zha_data
-from homeassistant.core import HomeAssistant
+from homeassistant.const import (
+    EVENT_HOMEASSISTANT_STOP,
+    MAJOR_VERSION,
+    MINOR_VERSION,
+    Platform,
+)
+from homeassistant.core import CoreState, HomeAssistant
 from homeassistant.helpers.event import async_call_later
 from homeassistant.setup import async_setup_component
 
@@ -203,3 +209,26 @@ async def test_zha_retry_unique_ids(
     await hass.config_entries.async_unload(config_entry.entry_id)
 
     assert "does not generate unique IDs" not in caplog.text
+
+
+async def test_shutdown_on_ha_stop(
+    hass: HomeAssistant,
+    config_entry: MockConfigEntry,
+    mock_zigpy_connect: ControllerApplication,
+) -> None:
+    """Test that the ZHA gateway is stopped when HA is shut down."""
+    config_entry.add_to_hass(hass)
+
+    await hass.config_entries.async_setup(config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    zha_data = get_zha_data(hass)
+
+    with patch.object(
+        zha_data.gateway, "shutdown", wraps=zha_data.gateway.shutdown
+    ) as mock_shutdown:
+        hass.bus.async_fire(EVENT_HOMEASSISTANT_STOP)
+        hass.state = CoreState.stopping
+        await hass.async_block_till_done()
+
+    assert len(mock_shutdown.mock_calls) == 1