Update ZHA entity discovery tests (#31191)

* ZHA registry test fixtures refactoring.

* Update ZHA test devices list.

* Use channel.id for event relay channels test.

* Add zigpy_device factory fixture.

* Add node descriptor to mock zigpy devices.

* Use node descriptor for ZHA device tests.

* Update ZHA discovery test.

check for unique_id and channels passed into each entity.

* Address comments.

tests/components/zha/common.py

@@ -78,7 +78,7 @@ def patch_cluster(cluster):
 class FakeDevice:
     """Fake device for mocking zigpy."""
 
-    def __init__(self, ieee, manufacturer, model):
+    def __init__(self, ieee, manufacturer, model, node_desc=None):
         """Init fake device."""
         self._application = APPLICATION
         self.ieee = zigpy.types.EUI64.convert(ieee)
@@ -95,6 +95,9 @@ class FakeDevice:
         self.node_desc = zigpy.zdo.types.NodeDescriptor()
         self.add_to_group = CoroutineMock()
         self.remove_from_group = CoroutineMock()
+        if node_desc is None:
+            node_desc = b"\x02@\x807\x10\x7fd\x00\x00*d\x00\x00"
+        self.node_desc = zigpy.zdo.types.NodeDescriptor.deserialize(node_desc)[0]
 
 
 def make_device(endpoints, ieee, manufacturer, model):

tests/components/zha/conftest.py

@@ -2,6 +2,7 @@
 from unittest import mock
 from unittest.mock import patch
 
+import asynctest
 import pytest
 import zigpy
 from zigpy.application import ControllerApplication
@@ -12,7 +13,7 @@ from homeassistant.components.zha.core.gateway import ZHAGateway
 from homeassistant.components.zha.core.store import async_get_registry
 from homeassistant.helpers.device_registry import async_get_registry as get_dev_reg
 
-from .common import async_setup_entry
+from .common import FakeDevice, FakeEndpoint, async_setup_entry
 
 FIXTURE_GRP_ID = 0x1001
 FIXTURE_GRP_NAME = "fixture group"
@@ -70,3 +71,52 @@ async def setup_zha(hass, config_entry):
         # init ZHA
         await hass.config_entries.async_forward_entry_setup(config_entry, DOMAIN)
         await hass.async_block_till_done()
+
+
+@pytest.fixture
+def channel():
+    """Channel mock factory fixture."""
+
+    def channel(name: str, cluster_id: int, endpoint_id: int = 1):
+        ch = mock.MagicMock()
+        ch.name = name
+        ch.generic_id = f"channel_0x{cluster_id:04x}"
+        ch.id = f"{endpoint_id}:0x{cluster_id:04x}"
+        ch.async_configure = asynctest.CoroutineMock()
+        ch.async_initialize = asynctest.CoroutineMock()
+        return ch
+
+    return channel
+
+
+@pytest.fixture
+def zigpy_device_mock():
+    """Make a fake device using the specified cluster classes."""
+
+    def _mock_dev(
+        endpoints,
+        ieee="00:0d:6f:00:0a:90:69:e7",
+        manufacturer="FakeManufacturer",
+        model="FakeModel",
+        node_desc=b"\x02@\x807\x10\x7fd\x00\x00*d\x00\x00",
+    ):
+        """Make a fake device using the specified cluster classes."""
+        device = FakeDevice(ieee, manufacturer, model, node_desc)
+        for epid, ep in endpoints.items():
+            endpoint = FakeEndpoint(manufacturer, model, epid)
+            endpoint.device = device
+            device.endpoints[epid] = endpoint
+            endpoint.device_type = ep["device_type"]
+            profile_id = ep.get("profile_id")
+            if profile_id:
+                endpoint.profile_id = profile_id
+
+            for cluster_id in ep.get("in_clusters", []):
+                endpoint.add_input_cluster(cluster_id)
+
+            for cluster_id in ep.get("out_clusters", []):
+                endpoint.add_output_cluster(cluster_id)
+
+        return device
+
+    return _mock_dev

tests/components/zha/test_channels.py

@@ -6,8 +6,6 @@ import homeassistant.components.zha.core.channels as channels
 import homeassistant.components.zha.core.device as zha_device
 import homeassistant.components.zha.core.registries as registries
 
-from .common import make_device
-
 
 @pytest.fixture
 def ieee():
@@ -64,9 +62,11 @@ def nwk():
         (0x1000, 1, {}),
     ],
 )
-async def test_in_channel_config(cluster_id, bind_count, attrs, zha_gateway, hass):
+async def test_in_channel_config(
+    cluster_id, bind_count, attrs, zha_gateway, hass, zigpy_device_mock
+):
     """Test ZHA core channel configuration for input clusters."""
-    zigpy_dev = make_device(
+    zigpy_dev = zigpy_device_mock(
         {1: {"in_clusters": [cluster_id], "out_clusters": [], "device_type": 0x1234}},
         "00:11:22:33:44:55:66:77",
         "test manufacturer",
@@ -120,9 +120,11 @@ async def test_in_channel_config(cluster_id, bind_count, attrs, zha_gateway, has
         (0x1000, 1),
     ],
 )
-async def test_out_channel_config(cluster_id, bind_count, zha_gateway, hass):
+async def test_out_channel_config(
+    cluster_id, bind_count, zha_gateway, hass, zigpy_device_mock
+):
     """Test ZHA core channel configuration for output clusters."""
-    zigpy_dev = make_device(
+    zigpy_dev = zigpy_device_mock(
         {1: {"out_clusters": [cluster_id], "in_clusters": [], "device_type": 0x1234}},
         "00:11:22:33:44:55:66:77",
         "test manufacturer",

tests/components/zha/test_discover.py

@@ -1,34 +1,50 @@
 """Test zha device discovery."""
 
 import asyncio
+import re
 from unittest import mock
 
 import pytest
 
-from homeassistant.components.zha.core.channels import EventRelayChannel
 import homeassistant.components.zha.core.const as zha_const
 import homeassistant.components.zha.core.discovery as disc
 import homeassistant.components.zha.core.gateway as core_zha_gw
+import homeassistant.helpers.entity_registry
 
-from .common import make_device
 from .zha_devices_list import DEVICES
 
+NO_TAIL_ID = re.compile("_\\d$")
+
 
 @pytest.mark.parametrize("device", DEVICES)
-async def test_devices(device, zha_gateway: core_zha_gw.ZHAGateway, hass, config_entry):
+async def test_devices(
+    device,
+    zha_gateway: core_zha_gw.ZHAGateway,
+    hass,
+    config_entry,
+    zigpy_device_mock,
+    monkeypatch,
+):
     """Test device discovery."""
 
-    zigpy_device = make_device(
+    zigpy_device = zigpy_device_mock(
         device["endpoints"],
         "00:11:22:33:44:55:66:77",
         device["manufacturer"],
         device["model"],
+        node_desc=device["node_descriptor"],
+    )
+
+    _dispatch = mock.MagicMock(wraps=disc.async_dispatch_discovery_info)
+    monkeypatch.setattr(core_zha_gw, "async_dispatch_discovery_info", _dispatch)
+    entity_registry = await homeassistant.helpers.entity_registry.async_get_registry(
+        hass
     )
 
     with mock.patch(
         "homeassistant.components.zha.core.discovery._async_create_cluster_channel",
         wraps=disc._async_create_cluster_channel,
-    ) as cr_ch:
+    ):
         await zha_gateway.async_device_restored(zigpy_device)
         await hass.async_block_till_done()
         tasks = [
@@ -45,11 +61,25 @@ async def test_devices(device, zha_gateway: core_zha_gw.ZHAGateway, hass, config
             ent for ent in entity_ids if ent.split(".")[0] in zha_const.COMPONENTS
         }
 
-        event_channels = {
-            arg[0].cluster_id
-            for arg, kwarg in cr_ch.call_args_list
-            if kwarg.get("channel_class") == EventRelayChannel
+        zha_dev = zha_gateway.get_device(zigpy_device.ieee)
+        event_channels = {  # pylint: disable=protected-access
+            ch.id for ch in zha_dev._relay_channels.values()
         }
 
         assert zha_entities == set(device["entities"])
         assert event_channels == set(device["event_channels"])
+
+        entity_map = device["entity_map"]
+        for calls in _dispatch.call_args_list:
+            discovery_info = calls[0][2]
+            unique_id = discovery_info["unique_id"]
+            channels = discovery_info["channels"]
+            component = discovery_info["component"]
+            key = (component, unique_id)
+            entity_id = entity_registry.async_get_entity_id(component, "zha", unique_id)
+
+            assert key in entity_map
+            assert entity_id is not None
+            no_tail_id = NO_TAIL_ID.sub("", entity_map[key]["entity_id"])
+            assert entity_id.startswith(no_tail_id)
+            assert set([ch.name for ch in channels]) == set(entity_map[key]["channels"])

tests/components/zha/test_registries.py

@@ -19,16 +19,10 @@ def zha_device():
 
 
 @pytest.fixture
-def channels():
+def channels(channel):
     """Return a mock of channels."""
 
-    def channel(name, chan_id):
-        ch = mock.MagicMock()
-        ch.name = name
-        ch.generic_id = chan_id
-        return ch
-
-    return [channel("level", "channel_0x0008"), channel("on_off", "channel_0x0006")]
+    return [channel("level", 8), channel("on_off", 6)]
 
 
 @pytest.mark.parametrize(