Migrate ZHA config entries to derive unique_id from the Zigbee EPID (home-assistant#154489)

Author: puddly

homeassistant/components/zha/__init__.py

@@ -48,6 +48,7 @@
     HAZHAData,
     ZHAGatewayProxy,
     create_zha_config,
+    get_config_entry_unique_id,
     get_zha_data,
 )
 from .radio_manager import ZhaRadioManager
@@ -198,6 +199,11 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
 
     repairs.async_delete_blocking_issues(hass)
 
+    # Set unique_id if it was not migrated previously
+    if not config_entry.unique_id or not config_entry.unique_id.startswith("epid="):
+        unique_id = get_config_entry_unique_id(zha_gateway.state.network_info)
+        hass.config_entries.async_update_entry(config_entry, unique_id=unique_id)
+
     ha_zha_data.gateway_proxy = ZHAGatewayProxy(hass, config_entry, zha_gateway)
 
     manufacturer = zha_gateway.state.node_info.manufacturer
@@ -313,5 +319,26 @@ async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
 
         hass.config_entries.async_update_entry(config_entry, data=data, version=4)
 
+    if config_entry.version == 4:
+        radio_mgr = ZhaRadioManager.from_config_entry(hass, config_entry)
+        await radio_mgr.async_read_backups_from_database()
+
+        if radio_mgr.backups:
+            # We migrate all ZHA config entries to use a `unique_id` specific to the
+            # Zigbee network, not to the hardware
+            backup = radio_mgr.backups[0]
+            hass.config_entries.async_update_entry(
+                config_entry,
+                unique_id=get_config_entry_unique_id(backup.network_info),
+                version=5,
+            )
+        else:
+            # If no backups are available, the unique_id will be set when the network is
+            # loaded during setup
+            hass.config_entries.async_update_entry(
+                config_entry,
+                version=5,
+            )
+
     _LOGGER.info("Migration to version %s successful", config_entry.version)
     return True

homeassistant/components/zha/config_flow.py

@@ -47,7 +47,7 @@
 from homeassistant.util import dt as dt_util
 
 from .const import CONF_BAUDRATE, CONF_FLOW_CONTROL, CONF_RADIO_TYPE, DOMAIN
-from .helpers import get_zha_gateway
+from .helpers import get_config_entry_unique_id, get_zha_gateway
 from .radio_manager import (
     DEVICE_SCHEMA,
     HARDWARE_DISCOVERY_SCHEMA,
@@ -544,6 +544,8 @@ async def async_step_form_new_network(
     ) -> ConfigFlowResult:
         """Form a brand-new network."""
         await self._radio_mgr.async_form_network()
+        # Load the newly formed network settings to get the network info
+        await self._radio_mgr.async_load_network_settings()
         return await self._async_create_radio_entry()
 
     def _parse_uploaded_backup(
@@ -668,7 +670,7 @@ async def async_step_maybe_confirm_ezsp_restore(
 class ZhaConfigFlowHandler(BaseZhaFlow, ConfigFlow, domain=DOMAIN):
     """Handle a config flow."""
 
-    VERSION = 4
+    VERSION = 5
 
     async def _set_unique_id_and_update_ignored_flow(
         self, unique_id: str, device_path: str
@@ -927,6 +929,15 @@ async def _async_create_radio_entry(self) -> ConfigFlowResult:
                 reason="reconfigure_successful",
             )
         if not zha_config_entries:
+            # Load network settings from the radio to get the EPID
+            await self._radio_mgr.async_load_network_settings()
+            assert self._radio_mgr.current_settings is not None
+
+            unique_id = get_config_entry_unique_id(
+                self._radio_mgr.current_settings.network_info
+            )
+            await self.async_set_unique_id(unique_id)
+
             return self.async_create_entry(
                 title=self._title,
                 data=data,

homeassistant/components/zha/helpers.py

@@ -90,6 +90,7 @@
 )
 import zigpy.exceptions
 from zigpy.profiles import PROFILES
+from zigpy.state import NetworkInfo
 import zigpy.types
 from zigpy.types import EUI64
 import zigpy.util
@@ -1390,3 +1391,8 @@ async def handler(self: _EntityT, *args: _P.args, **kwargs: _P.kwargs) -> None:
 def exclude_none_values(obj: Mapping[str, Any]) -> dict[str, Any]:
     """Return a new dictionary excluding keys with None values."""
     return {k: v for k, v in obj.items() if v is not None}
+
+
+def get_config_entry_unique_id(network_info: NetworkInfo) -> str:
+    """Generate a unique id for a config entry based on the network info."""
+    return f"epid={network_info.extended_pan_id}".lower()

tests/components/zha/conftest.py

@@ -195,7 +195,7 @@ async def zigpy_app_controller():
 async def config_entry_fixture() -> MockConfigEntry:
     """Fixture representing a config entry."""
     return MockConfigEntry(
-        version=4,
+        version=5,
         domain=zha_const.DOMAIN,
         data={
             zigpy.config.CONF_DEVICE: {

tests/components/zha/snapshots/test_diagnostics.ambr

@@ -117,8 +117,8 @@
       'subentries': list([
       ]),
       'title': 'Mock Title',
-      'unique_id': None,
-      'version': 4,
+      'unique_id': '**REDACTED**',
+      'version': 5,
     }),
     'devices': list([
       dict({

tests/components/zha/test_config_flow.py

@@ -96,6 +96,9 @@ def mock_app() -> Generator[AsyncMock]:
     mock_app = AsyncMock()
     mock_app.backups = create_autospec(BackupManager, instance=True)
     mock_app.backups.backups = []
+    mock_app.state.network_info.extended_pan_id = zigpy.types.EUI64.convert(
+        "AABBCCDDEE000000"
+    )
     mock_app.state.network_info.metadata = {
         "ezsp": {
             "can_burn_userdata_custom_eui64": True,
@@ -175,7 +178,7 @@ def com_port(device="/dev/ttyUSB1234") -> ListPortInfo:
         (
             # TubesZB, old ESPHome devices (ZNP)
             "tubeszb-cc2652-poe",
-            "tubeszb-cc2652-poe",
+            "epid=aa:bb:cc:dd:ee:00:00:00",
             RadioType.znp,
             ZeroconfServiceInfo(
                 ip_address=ip_address("192.168.1.200"),
@@ -198,7 +201,7 @@ def com_port(device="/dev/ttyUSB1234") -> ListPortInfo:
         (
             # TubesZB, old ESPHome device (EFR32)
             "tubeszb-efr32-poe",
-            "tubeszb-efr32-poe",
+            "epid=aa:bb:cc:dd:ee:00:00:00",
             RadioType.ezsp,
             ZeroconfServiceInfo(
                 ip_address=ip_address("192.168.1.200"),
@@ -221,7 +224,7 @@ def com_port(device="/dev/ttyUSB1234") -> ListPortInfo:
         (
             # TubesZB, newer devices
             "TubeZB",
-            "tubeszb-cc2652-poe",
+            "epid=aa:bb:cc:dd:ee:00:00:00",
             RadioType.znp,
             ZeroconfServiceInfo(
                 ip_address=ip_address("192.168.1.200"),
@@ -242,7 +245,7 @@ def com_port(device="/dev/ttyUSB1234") -> ListPortInfo:
         (
             # Expected format for all new devices
             "Some Zigbee Gateway (12345)",
-            "aabbccddeeff",
+            "epid=aa:bb:cc:dd:ee:00:00:00",
             RadioType.znp,
             ZeroconfServiceInfo(
                 ip_address=ip_address("192.168.1.200"),
@@ -1627,8 +1630,15 @@ async def test_formation_strategy_form_initial_network(
     advanced_pick_radio: RadioPicker, mock_app: AsyncMock, hass: HomeAssistant
 ) -> None:
     """Test forming a new network, with no previous settings on the radio."""
+    # Initially, no network is formed
     mock_app.load_network_info = AsyncMock(side_effect=NetworkNotFormed())
 
+    # After form_network is called, load_network_info should return the network settings
+    async def form_network_side_effect(*args, **kwargs):
+        mock_app.load_network_info = AsyncMock(return_value=mock_app.state.network_info)
+
+    mock_app.form_network.side_effect = form_network_side_effect
+
     result = await advanced_pick_radio(RadioType.ezsp)
     result2 = await hass.config_entries.flow.async_configure(
         result["flow_id"],
@@ -1648,8 +1658,16 @@ async def test_onboarding_auto_formation_new_hardware(
     mock_app: AsyncMock, hass: HomeAssistant
 ) -> None:
     """Test auto network formation with new hardware during onboarding."""
+    # Initially, no network is formed
     mock_app.load_network_info = AsyncMock(side_effect=NetworkNotFormed())
     mock_app.get_device = MagicMock(return_value=MagicMock(spec=zigpy.device.Device))
+
+    # After form_network is called, load_network_info should return the network settings
+    async def form_network_side_effect(*args, **kwargs):
+        mock_app.load_network_info = AsyncMock(return_value=mock_app.state.network_info)
+
+    mock_app.form_network.side_effect = form_network_side_effect
+
     discovery_info = UsbServiceInfo(
         device="/dev/ttyZIGBEE",
         pid="AAAA",

tests/components/zha/test_homeassistant_hardware.py

@@ -35,7 +35,7 @@ async def test_get_firmware_info_normal(hass: HomeAssistant) -> None:
             },
             "radio_type": "ezsp",
         },
-        version=4,
+        version=5,
     )
     zha.add_to_hass(hass)
     zha.mock_state(hass, ConfigEntryState.LOADED)
@@ -87,7 +87,7 @@ async def test_get_firmware_info_errors(
         domain="zha",
         unique_id="some_unique_id",
         data=data,
-        version=4,
+        version=5,
     )
     zha.add_to_hass(hass)
 

tests/components/zha/test_init.py

@@ -73,7 +73,7 @@ async def test_migration_from_v1_no_baudrate(
     assert CONF_DEVICE in config_entry_v1.data
     assert config_entry_v1.data[CONF_DEVICE][CONF_DEVICE_PATH] == DATA_PORT_PATH
     assert CONF_USB_PATH not in config_entry_v1.data
-    assert config_entry_v1.version == 4
+    assert config_entry_v1.version == 5
 
 
 @patch("homeassistant.components.zha.async_setup_entry", AsyncMock(return_value=True))
@@ -90,7 +90,7 @@ async def test_migration_from_v1_with_baudrate(
     assert CONF_USB_PATH not in config_entry_v1.data
     assert CONF_BAUDRATE in config_entry_v1.data[CONF_DEVICE]
     assert config_entry_v1.data[CONF_DEVICE][CONF_BAUDRATE] == 115200
-    assert config_entry_v1.version == 4
+    assert config_entry_v1.version == 5
 
 
 @patch("homeassistant.components.zha.async_setup_entry", AsyncMock(return_value=True))
@@ -105,7 +105,7 @@ async def test_migration_from_v1_wrong_baudrate(
     assert CONF_DEVICE in config_entry_v1.data
     assert config_entry_v1.data[CONF_DEVICE][CONF_DEVICE_PATH] == DATA_PORT_PATH
     assert CONF_USB_PATH not in config_entry_v1.data
-    assert config_entry_v1.version == 4
+    assert config_entry_v1.version == 5
 
 
 @pytest.mark.skipif(
@@ -167,7 +167,7 @@ async def test_setup_with_v3_cleaning_uri(
                 CONF_FLOW_CONTROL: None,
             },
         },
-        version=4,
+        version=5,
     )
     config_entry_v4.add_to_hass(hass)
 
@@ -177,7 +177,7 @@ async def test_setup_with_v3_cleaning_uri(
 
     assert config_entry_v4.data[CONF_RADIO_TYPE] == DATA_RADIO_TYPE
     assert config_entry_v4.data[CONF_DEVICE][CONF_DEVICE_PATH] == cleaned_path
-    assert config_entry_v4.version == 4
+    assert config_entry_v4.version == 5
 
 
 @pytest.mark.parametrize(