Fix ZHA startup creating entities with non-unique IDs (#99679)

* Make the ZHAGateway initialization restartable so entities are unique

* Add a unit test

homeassistant/components/zha/__init__.py

@@ -134,7 +134,9 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
     else:
         _LOGGER.debug("ZHA storage file does not exist or was already removed")
 
-    zha_gateway = ZHAGateway(hass, config, config_entry)
+    # Re-use the gateway object between ZHA reloads
+    if (zha_gateway := zha_data.get(DATA_ZHA_GATEWAY)) is None:
+        zha_gateway = ZHAGateway(hass, config, config_entry)
 
     try:
         await zha_gateway.async_initialize()

homeassistant/components/zha/core/gateway.py

@@ -149,6 +149,12 @@ class ZHAGateway:
         self.config_entry = config_entry
         self._unsubs: list[Callable[[], None]] = []
 
+        discovery.PROBE.initialize(self._hass)
+        discovery.GROUP_PROBE.initialize(self._hass)
+
+        self.ha_device_registry = dr.async_get(self._hass)
+        self.ha_entity_registry = er.async_get(self._hass)
+
     def get_application_controller_data(self) -> tuple[ControllerApplication, dict]:
         """Get an uninitialized instance of a zigpy `ControllerApplication`."""
         radio_type = self.config_entry.data[CONF_RADIO_TYPE]
@@ -191,12 +197,6 @@ class ZHAGateway:
 
     async def async_initialize(self) -> None:
         """Initialize controller and connect radio."""
-        discovery.PROBE.initialize(self._hass)
-        discovery.GROUP_PROBE.initialize(self._hass)
-
-        self.ha_device_registry = dr.async_get(self._hass)
-        self.ha_entity_registry = er.async_get(self._hass)
-
         app_controller_cls, app_config = self.get_application_controller_data()
         self.application_controller = await app_controller_cls.new(
             config=app_config,

tests/components/zha/test_init.py

@@ -1,8 +1,10 @@
 """Tests for ZHA integration init."""
+import asyncio
 from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
 from zigpy.config import CONF_DEVICE, CONF_DEVICE_PATH
+from zigpy.exceptions import TransientConnectionError
 
 from homeassistant.components.zha import async_setup_entry
 from homeassistant.components.zha.core.const import (
@@ -11,10 +13,13 @@ from homeassistant.components.zha.core.const import (
     CONF_USB_PATH,
     DOMAIN,
 )
-from homeassistant.const import MAJOR_VERSION, MINOR_VERSION
+from homeassistant.const import MAJOR_VERSION, MINOR_VERSION, Platform
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers.event import async_call_later
 from homeassistant.setup import async_setup_component
 
+from .test_light import LIGHT_ON_OFF
+
 from tests.common import MockConfigEntry
 
 DATA_RADIO_TYPE = "deconz"
@@ -157,3 +162,46 @@ async def test_setup_with_v3_cleaning_uri(
     assert config_entry_v3.data[CONF_RADIO_TYPE] == DATA_RADIO_TYPE
     assert config_entry_v3.data[CONF_DEVICE][CONF_DEVICE_PATH] == cleaned_path
     assert config_entry_v3.version == 3
+
+
+@patch(
+    "homeassistant.components.zha.PLATFORMS",
+    [Platform.LIGHT, Platform.BUTTON, Platform.SENSOR, Platform.SELECT],
+)
+async def test_zha_retry_unique_ids(
+    hass: HomeAssistant,
+    config_entry: MockConfigEntry,
+    zigpy_device_mock,
+    mock_zigpy_connect,
+    caplog,
+) -> None:
+    """Test that ZHA retrying creates unique entity IDs."""
+
+    config_entry.add_to_hass(hass)
+
+    # Ensure we have some device to try to load
+    app = mock_zigpy_connect.return_value
+    light = zigpy_device_mock(LIGHT_ON_OFF)
+    app.devices[light.ieee] = light
+
+    # Re-try setup but have it fail once, so entities have two chances to be created
+    with patch.object(
+        app,
+        "startup",
+        side_effect=[TransientConnectionError(), None],
+    ) as mock_connect:
+        with patch(
+            "homeassistant.config_entries.async_call_later",
+            lambda hass, delay, action: async_call_later(hass, 0, action),
+        ):
+            await hass.config_entries.async_setup(config_entry.entry_id)
+            await hass.async_block_till_done()
+
+            # Wait for the config entry setup to retry
+            await asyncio.sleep(0.1)
+
+        assert len(mock_connect.mock_calls) == 2
+
+    await hass.config_entries.async_unload(config_entry.entry_id)
+
+    assert "does not generate unique IDs" not in caplog.text