Adapt zha to set via_device_id in DeviceInfo (#178168)

homeassistant/components/zha/entity.py

@@ -12,14 +12,9 @@ from propcache.api import cached_property
 from zha.application.platforms import EntityStateChangedEvent
 from zha.mixins import LogMixin
 
-from homeassistant.const import (
-    ATTR_MANUFACTURER,
-    ATTR_MODEL,
-    ATTR_NAME,
-    ATTR_VIA_DEVICE,
-    EntityCategory,
-)
+from homeassistant.const import ATTR_MANUFACTURER, ATTR_MODEL, ATTR_NAME, EntityCategory
 from homeassistant.core import State, callback
+from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.device_registry import CONNECTION_ZIGBEE, DeviceInfo
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity import Entity
@@ -134,7 +129,8 @@ class ZHAEntity(LogMixin, RestoreEntity, Entity):
         """Return a device description for device registry."""
         zha_device_info = self.entity_data.device_proxy.device_info
         ieee = zha_device_info["ieee"]
-        zha_gateway = self.entity_data.device_proxy.gateway_proxy.gateway
+        gateway_proxy = self.entity_data.device_proxy.gateway_proxy
+        zha_gateway = gateway_proxy.gateway
 
         device_info = DeviceInfo(
             connections={(CONNECTION_ZIGBEE, ieee)},
@@ -143,10 +139,14 @@ class ZHAEntity(LogMixin, RestoreEntity, Entity):
             model=zha_device_info[ATTR_MODEL],
             name=zha_device_info[ATTR_NAME],
         )
-        if ieee != str(zha_gateway.state.node_info.ieee):
-            device_info[ATTR_VIA_DEVICE] = (
-                DOMAIN,
-                str(zha_gateway.state.node_info.ieee),
+        coordinator_ieee = str(zha_gateway.state.node_info.ieee)
+        if ieee != coordinator_ieee:
+            # The coordinator device is registered before platforms are set up,
+            # so it is always present when a child entity's device_info is built.
+            device_info["via_device_id"] = dr.async_get_device_id_by_identifier(
+                gateway_proxy.hass,
+                (DOMAIN, coordinator_ieee),
+                config_entry_id=gateway_proxy.config_entry.entry_id,
             )
         return device_info
 

tests/components/zha/test_entity.py

@@ -6,20 +6,24 @@ from zigpy.device import Device
 from zigpy.profiles import zha
 from zigpy.zcl.clusters import general
 
+from homeassistant.components.zha.const import DOMAIN
 from homeassistant.components.zha.helpers import get_zha_gateway
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr
 
 from .conftest import SIG_EP_INPUT, SIG_EP_OUTPUT, SIG_EP_PROFILE, SIG_EP_TYPE
 
+from tests.common import MockConfigEntry
+
 
 async def test_device_registry_via_device(
     hass: HomeAssistant,
+    config_entry: MockConfigEntry,
     setup_zha: Callable[..., Coroutine[None]],
     zigpy_device_mock: Callable[..., Device],
     device_registry: dr.DeviceRegistry,
 ) -> None:
-    """Test ZHA `via_device` is set correctly."""
+    """Test a ZHA device links to the coordinator device via via_device_id."""
 
     await setup_zha()
     gateway = get_zha_gateway(hass)
@@ -39,12 +43,13 @@ async def test_device_registry_via_device(
     await gateway.async_device_initialized(zigpy_device)
     await hass.async_block_till_done(wait_background_tasks=True)
 
-    reg_coordinator_device = device_registry.async_get_device(
-        identifiers={("zha", str(gateway.state.node_info.ieee))}
+    coordinator_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, str(gateway.state.node_info.ieee)), config_entry.entry_id
     )
+    assert coordinator_device is not None
 
-    reg_device = device_registry.async_get_device(
-        identifiers={("zha", str(zha_device.ieee))}
+    reg_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, str(zha_device.ieee)), config_entry.entry_id
     )
-
-    assert reg_device.via_device_id == reg_coordinator_device.id
+    assert reg_device is not None
+    assert reg_device.via_device_id == coordinator_device.id