Group by endpoints and not by devices for ZHA Zigbee groups (#34583)

* start implementation
* handle members correctly
* fix group member info
* align groupable devices with group members
* handle group endpoint adding and removing
* update add group
* update group and group member
* update create group
* remove domain check
* update test
* remove temporary 2nd groupable device api
* update test
* rename validator - review comment
* fix test that was never running
* additional testing
* fix coordinator descriptors
* remove check that was done twice
* update test
* Use AsyncMock()
Co-authored-by: Alexei Chetroi <lexoid@gmail.com>

homeassistant/components/zha/api.py

@@ -53,6 +53,7 @@ from .core.const import (
     WARNING_DEVICE_STROBE_HIGH,
     WARNING_DEVICE_STROBE_YES,
 )
+from .core.group import GroupMember
 from .core.helpers import async_is_bindable_target, get_matched_clusters
 
 _LOGGER = logging.getLogger(__name__)
@@ -209,7 +210,7 @@ async def websocket_get_devices(hass, connection, msg):
     """Get ZHA devices."""
     zha_gateway = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY]
 
-    devices = [device.async_get_info() for device in zha_gateway.devices.values()]
+    devices = [device.zha_device_info for device in zha_gateway.devices.values()]
 
     connection.send_result(msg[ID], devices)
 
@@ -221,13 +222,35 @@ async def websocket_get_groupable_devices(hass, connection, msg):
     """Get ZHA devices that can be grouped."""
     zha_gateway = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY]
 
-    devices = [
+    devices = [device for device in zha_gateway.devices.values() if device.is_groupable]
-        device.async_get_info()
+    groupable_devices = []
-        for device in zha_gateway.devices.values()
-        if device.is_groupable or device.is_coordinator
-    ]
 
-    connection.send_result(msg[ID], devices)
+    for device in devices:
+        entity_refs = zha_gateway.device_registry.get(device.ieee)
+        for ep_id in device.async_get_groupable_endpoints():
+            groupable_devices.append(
+                {
+                    "endpoint_id": ep_id,
+                    "entities": [
+                        {
+                            "name": zha_gateway.ha_entity_registry.async_get(
+                                entity_ref.reference_id
+                            ).name,
+                            "original_name": zha_gateway.ha_entity_registry.async_get(
+                                entity_ref.reference_id
+                            ).original_name,
+                        }
+                        for entity_ref in entity_refs
+                        if list(entity_ref.cluster_channels.values())[
+                            0
+                        ].cluster.endpoint.endpoint_id
+                        == ep_id
+                    ],
+                    "device": device.zha_device_info,
+                }
+            )
+
+    connection.send_result(msg[ID], groupable_devices)
 
 
 @websocket_api.require_admin
@@ -236,7 +259,7 @@ async def websocket_get_groupable_devices(hass, connection, msg):
 async def websocket_get_groups(hass, connection, msg):
     """Get ZHA groups."""
     zha_gateway = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY]
-    groups = [group.async_get_info() for group in zha_gateway.groups.values()]
+    groups = [group.group_info for group in zha_gateway.groups.values()]
     connection.send_result(msg[ID], groups)
 
 
@@ -251,7 +274,7 @@ async def websocket_get_device(hass, connection, msg):
     ieee = msg[ATTR_IEEE]
     device = None
     if ieee in zha_gateway.devices:
-        device = zha_gateway.devices[ieee].async_get_info()
+        device = zha_gateway.devices[ieee].zha_device_info
     if not device:
         connection.send_message(
             websocket_api.error_message(
@@ -274,7 +297,7 @@ async def websocket_get_group(hass, connection, msg):
     group = None
 
     if group_id in zha_gateway.groups:
-        group = zha_gateway.groups.get(group_id).async_get_info()
+        group = zha_gateway.groups.get(group_id).group_info
     if not group:
         connection.send_message(
             websocket_api.error_message(
@@ -285,13 +308,27 @@ async def websocket_get_group(hass, connection, msg):
     connection.send_result(msg[ID], group)
 
 
+def cv_group_member(value: Any) -> GroupMember:
+    """Validate and transform a group member."""
+    if not isinstance(value, Mapping):
+        raise vol.Invalid("Not a group member")
+    try:
+        group_member = GroupMember(
+            ieee=EUI64.convert(value["ieee"]), endpoint_id=value["endpoint_id"]
+        )
+    except KeyError:
+        raise vol.Invalid("Not a group member")
+
+    return group_member
+
+
 @websocket_api.require_admin
 @websocket_api.async_response
 @websocket_api.websocket_command(
     {
         vol.Required(TYPE): "zha/group/add",
         vol.Required(GROUP_NAME): cv.string,
-        vol.Optional(ATTR_MEMBERS): vol.All(cv.ensure_list, [EUI64.convert]),
+        vol.Optional(ATTR_MEMBERS): vol.All(cv.ensure_list, [cv_group_member]),
     }
 )
 async def websocket_add_group(hass, connection, msg):
@@ -300,7 +337,7 @@ async def websocket_add_group(hass, connection, msg):
     group_name = msg[GROUP_NAME]
     members = msg.get(ATTR_MEMBERS)
     group = await zha_gateway.async_create_zigpy_group(group_name, members)
-    connection.send_result(msg[ID], group.async_get_info())
+    connection.send_result(msg[ID], group.group_info)
 
 
 @websocket_api.require_admin
@@ -323,7 +360,7 @@ async def websocket_remove_groups(hass, connection, msg):
         await asyncio.gather(*tasks)
     else:
         await zha_gateway.async_remove_zigpy_group(group_ids[0])
-    ret_groups = [group.async_get_info() for group in zha_gateway.groups.values()]
+    ret_groups = [group.group_info for group in zha_gateway.groups.values()]
     connection.send_result(msg[ID], ret_groups)
 
 
@@ -333,7 +370,7 @@ async def websocket_remove_groups(hass, connection, msg):
     {
         vol.Required(TYPE): "zha/group/members/add",
         vol.Required(GROUP_ID): cv.positive_int,
-        vol.Required(ATTR_MEMBERS): vol.All(cv.ensure_list, [EUI64.convert]),
+        vol.Required(ATTR_MEMBERS): vol.All(cv.ensure_list, [cv_group_member]),
     }
 )
 async def websocket_add_group_members(hass, connection, msg):
@@ -353,7 +390,7 @@ async def websocket_add_group_members(hass, connection, msg):
             )
         )
         return
-    ret_group = zha_group.async_get_info()
+    ret_group = zha_group.group_info
     connection.send_result(msg[ID], ret_group)
 
 
@@ -363,7 +400,7 @@ async def websocket_add_group_members(hass, connection, msg):
     {
         vol.Required(TYPE): "zha/group/members/remove",
         vol.Required(GROUP_ID): cv.positive_int,
-        vol.Required(ATTR_MEMBERS): vol.All(cv.ensure_list, [EUI64.convert]),
+        vol.Required(ATTR_MEMBERS): vol.All(cv.ensure_list, [cv_group_member]),
     }
 )
 async def websocket_remove_group_members(hass, connection, msg):
@@ -383,7 +420,7 @@ async def websocket_remove_group_members(hass, connection, msg):
             )
         )
         return
-    ret_group = zha_group.async_get_info()
+    ret_group = zha_group.group_info
     connection.send_result(msg[ID], ret_group)
 
 
@@ -608,7 +645,7 @@ async def websocket_get_bindable_devices(hass, connection, msg):
     source_device = zha_gateway.get_device(source_ieee)
 
     devices = [
-        device.async_get_info()
+        device.zha_device_info
         for device in zha_gateway.devices.values()
         if async_is_bindable_target(source_device, device)
     ]

homeassistant/components/zha/core/device.py

@@ -235,13 +235,9 @@ class ZHADevice(LogMixin):
     @property
     def is_groupable(self):
         """Return true if this device has a group cluster."""
-        if not self.available:
+        return self.is_coordinator or (
-            return False
+            self.available and self.async_get_groupable_endpoints()
-        clusters = self.async_get_clusters()
+        )
-        for cluster_map in clusters.values():
-            for clusters in cluster_map.values():
-                if Groups.cluster_id in clusters:
-                    return True
 
     @property
     def skip_configuration(self):
@@ -411,8 +407,8 @@ class ZHADevice(LogMixin):
         if self._zigpy_device.last_seen is None and last_seen is not None:
             self._zigpy_device.last_seen = last_seen
 
-    @callback
+    @property
-    def async_get_info(self):
+    def zha_device_info(self):
         """Get ZHA device information."""
         device_info = {}
         device_info.update(self.device_info)
@@ -442,6 +438,15 @@ class ZHADevice(LogMixin):
             if ep_id != 0
         }
 
+    @callback
+    def async_get_groupable_endpoints(self):
+        """Get device endpoints that have a group 'in' cluster."""
+        return [
+            ep_id
+            for (ep_id, clusters) in self.async_get_clusters().items()
+            if Groups.cluster_id in clusters[CLUSTER_TYPE_IN]
+        ]
+
     @callback
     def async_get_std_clusters(self):
         """Get ZHA and ZLL clusters for this device."""
@@ -557,7 +562,15 @@ class ZHADevice(LogMixin):
 
     async def async_add_to_group(self, group_id):
         """Add this device to the provided zigbee group."""
-        await self._zigpy_device.add_to_group(group_id)
+        try:
+            await self._zigpy_device.add_to_group(group_id)
+        except (zigpy.exceptions.ZigbeeException, asyncio.TimeoutError) as ex:
+            self.debug(
+                "Failed to add device '%s' to group: 0x%04x ex: %s",
+                self._zigpy_device.ieee,
+                group_id,
+                str(ex),
+            )
 
     async def async_remove_from_group(self, group_id):
         """Remove this device from the provided zigbee group."""
@@ -571,6 +584,34 @@ class ZHADevice(LogMixin):
                 str(ex),
             )
 
+    async def async_add_endpoint_to_group(self, endpoint_id, group_id):
+        """Add the device endpoint to the provided zigbee group."""
+        try:
+            await self._zigpy_device.endpoints[int(endpoint_id)].add_to_group(group_id)
+        except (zigpy.exceptions.ZigbeeException, asyncio.TimeoutError) as ex:
+            self.debug(
+                "Failed to add endpoint: %s for device: '%s' to group: 0x%04x ex: %s",
+                endpoint_id,
+                self._zigpy_device.ieee,
+                group_id,
+                str(ex),
+            )
+
+    async def async_remove_endpoint_from_group(self, endpoint_id, group_id):
+        """Remove the device endpoint from the provided zigbee group."""
+        try:
+            await self._zigpy_device.endpoints[int(endpoint_id)].remove_from_group(
+                group_id
+            )
+        except (zigpy.exceptions.ZigbeeException, asyncio.TimeoutError) as ex:
+            self.debug(
+                "Failed to remove endpoint: %s for device '%s' from group: 0x%04x ex: %s",
+                endpoint_id,
+                self._zigpy_device.ieee,
+                group_id,
+                str(ex),
+            )
+
     async def async_bind_to_group(self, group_id, cluster_bindings):
         """Directly bind this device to a group for the given clusters."""
         await self._async_group_binding_operation(

homeassistant/components/zha/core/discovery.py

@@ -235,21 +235,13 @@ class GroupProbe:
     ) -> List[str]:
         """Determine the entity domains for this group."""
         entity_domains: List[str] = []
-        if len(group.members) < 2:
-            _LOGGER.debug(
-                "Group: %s:0x%04x has less than 2 members so cannot default an entity domain",
-                group.name,
-                group.group_id,
-            )
-            return entity_domains
-
         zha_gateway = hass.data[zha_const.DATA_ZHA][zha_const.DATA_ZHA_GATEWAY]
         all_domain_occurrences = []
-        for device in group.members:
+        for member in group.members:
-            if device.is_coordinator:
+            if member.device.is_coordinator:
                 continue
             entities = async_entries_for_device(
-                zha_gateway.ha_entity_registry, device.device_id
+                zha_gateway.ha_entity_registry, member.device.device_id
             )
             all_domain_occurrences.extend(
                 [

homeassistant/components/zha/core/gateway.py

@@ -78,11 +78,11 @@ from .const import (
     ZHA_GW_RADIO_DESCRIPTION,
 )
 from .device import DeviceStatus, ZHADevice
-from .group import ZHAGroup
+from .group import GroupMember, ZHAGroup
 from .patches import apply_application_controller_patch
 from .registries import GROUP_ENTITY_DOMAINS, RADIO_TYPES
 from .store import async_get_registry
-from .typing import ZhaDeviceType, ZhaGroupType, ZigpyEndpointType, ZigpyGroupType
+from .typing import ZhaGroupType, ZigpyEndpointType, ZigpyGroupType
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -308,7 +308,7 @@ class ZHAGateway:
                 ZHA_GW_MSG,
                 {
                     ATTR_TYPE: gateway_message_type,
-                    ZHA_GW_MSG_GROUP_INFO: zha_group.async_get_info(),
+                    ZHA_GW_MSG_GROUP_INFO: zha_group.group_info,
                 },
             )
 
@@ -327,7 +327,7 @@ class ZHAGateway:
         zha_device = self._devices.pop(device.ieee, None)
         entity_refs = self._device_registry.pop(device.ieee, None)
         if zha_device is not None:
-            device_info = zha_device.async_get_info()
+            device_info = zha_device.zha_device_info
             zha_device.async_cleanup_handles()
             async_dispatcher_send(
                 self._hass, "{}_{}".format(SIGNAL_REMOVE, str(zha_device.ieee))
@@ -542,7 +542,7 @@ class ZHAGateway:
             )
             await self._async_device_joined(zha_device)
 
-        device_info = zha_device.async_get_info()
+        device_info = zha_device.zha_device_info
 
         async_dispatcher_send(
             self._hass,
@@ -571,11 +571,11 @@ class ZHAGateway:
         zha_device.update_available(True)
 
     async def async_create_zigpy_group(
-        self, name: str, members: List[ZhaDeviceType]
+        self, name: str, members: List[GroupMember]
     ) -> ZhaGroupType:
         """Create a new Zigpy Zigbee group."""
-        # we start with one to fill any gaps from a user removing existing groups
+        # we start with two to fill any gaps from a user removing existing groups
-        group_id = 1
+        group_id = 2
         while group_id in self.groups:
             group_id += 1
 
@@ -584,14 +584,19 @@ class ZHAGateway:
             self.application_controller.groups.add_group(group_id, name)
             if members is not None:
                 tasks = []
-                for ieee in members:
+                for member in members:
                     _LOGGER.debug(
-                        "Adding member with IEEE: %s to group: %s:0x%04x",
+                        "Adding member with IEEE: %s and endpoint id: %s to group: %s:0x%04x",
-                        ieee,
+                        member.ieee,
+                        member.endpoint_id,
                         name,
                         group_id,
                     )
-                    tasks.append(self.devices[ieee].async_add_to_group(group_id))
+                    tasks.append(
+                        self.devices[member.ieee].async_add_endpoint_to_group(
+                            member.endpoint_id, group_id
+                        )
+                    )
                 await asyncio.gather(*tasks)
         return self.groups.get(group_id)
 
@@ -604,7 +609,7 @@ class ZHAGateway:
         if group and group.members:
             tasks = []
             for member in group.members:
-                tasks.append(member.async_remove_from_group(group_id))
+                tasks.append(member.async_remove_from_group())
             if tasks:
                 await asyncio.gather(*tasks)
         self.application_controller.groups.pop(group_id)

homeassistant/components/zha/core/group.py

@@ -1,19 +1,110 @@
 """Group for Zigbee Home Automation."""
 import asyncio
+import collections
 import logging
 from typing import Any, Dict, List
 
-from zigpy.types.named import EUI64
+import zigpy.exceptions
 
-from homeassistant.core import callback
 from homeassistant.helpers.entity_registry import async_entries_for_device
 from homeassistant.helpers.typing import HomeAssistantType
 
 from .helpers import LogMixin
-from .typing import ZhaDeviceType, ZhaGatewayType, ZigpyEndpointType, ZigpyGroupType
+from .typing import (
+    ZhaDeviceType,
+    ZhaGatewayType,
+    ZhaGroupType,
+    ZigpyEndpointType,
+    ZigpyGroupType,
+)
 
 _LOGGER = logging.getLogger(__name__)
 
+GroupMember = collections.namedtuple("GroupMember", "ieee endpoint_id")
+GroupEntityReference = collections.namedtuple(
+    "GroupEntityReference", "name original_name entity_id"
+)
+
+
+class ZHAGroupMember(LogMixin):
+    """Composite object that represents a device endpoint in a Zigbee group."""
+
+    def __init__(
+        self, zha_group: ZhaGroupType, zha_device: ZhaDeviceType, endpoint_id: int
+    ):
+        """Initialize the group member."""
+        self._zha_group: ZhaGroupType = zha_group
+        self._zha_device: ZhaDeviceType = zha_device
+        self._endpoint_id: int = endpoint_id
+
+    @property
+    def group(self) -> ZhaGroupType:
+        """Return the group this member belongs to."""
+        return self._zha_group
+
+    @property
+    def endpoint_id(self) -> int:
+        """Return the endpoint id for this group member."""
+        return self._endpoint_id
+
+    @property
+    def endpoint(self) -> ZigpyEndpointType:
+        """Return the endpoint for this group member."""
+        return self._zha_device.device.endpoints.get(self.endpoint_id)
+
+    @property
+    def device(self) -> ZhaDeviceType:
+        """Return the zha device for this group member."""
+        return self._zha_device
+
+    @property
+    def member_info(self) -> Dict[str, Any]:
+        """Get ZHA group info."""
+        member_info: Dict[str, Any] = {}
+        member_info["endpoint_id"] = self.endpoint_id
+        member_info["device"] = self.device.zha_device_info
+        member_info["entities"] = self.associated_entities
+        return member_info
+
+    @property
+    def associated_entities(self) -> List[GroupEntityReference]:
+        """Return the list of entities that were derived from this endpoint."""
+        ha_entity_registry = self.device.gateway.ha_entity_registry
+        zha_device_registry = self.device.gateway.device_registry
+        return [
+            GroupEntityReference(
+                ha_entity_registry.async_get(entity_ref.reference_id).name,
+                ha_entity_registry.async_get(entity_ref.reference_id).original_name,
+                entity_ref.reference_id,
+            )._asdict()
+            for entity_ref in zha_device_registry.get(self.device.ieee)
+            if list(entity_ref.cluster_channels.values())[
+                0
+            ].cluster.endpoint.endpoint_id
+            == self.endpoint_id
+        ]
+
+    async def async_remove_from_group(self) -> None:
+        """Remove the device endpoint from the provided zigbee group."""
+        try:
+            await self._zha_device.device.endpoints[
+                self._endpoint_id
+            ].remove_from_group(self._zha_group.group_id)
+        except (zigpy.exceptions.ZigbeeException, asyncio.TimeoutError) as ex:
+            self.debug(
+                "Failed to remove endpoint: %s for device '%s' from group: 0x%04x ex: %s",
+                self._endpoint_id,
+                self._zha_device.ieee,
+                self._zha_group.group_id,
+                str(ex),
+            )
+
+    def log(self, level: int, msg: str, *args) -> None:
+        """Log a message."""
+        msg = f"[%s](%s): {msg}"
+        args = (f"0x{self._zha_group.group_id:04x}", self.endpoint_id) + args
+        _LOGGER.log(level, msg, *args)
+
 
 class ZHAGroup(LogMixin):
     """ZHA Zigbee group object."""
@@ -45,77 +136,79 @@ class ZHAGroup(LogMixin):
         return self._zigpy_group.endpoint
 
     @property
-    def members(self) -> List[ZhaDeviceType]:
+    def members(self) -> List[ZHAGroupMember]:
         """Return the ZHA devices that are members of this group."""
         return [
-            self._zha_gateway.devices.get(member_ieee[0])
+            ZHAGroupMember(
-            for member_ieee in self._zigpy_group.members.keys()
+                self, self._zha_gateway.devices.get(member_ieee), endpoint_id
-            if member_ieee[0] in self._zha_gateway.devices
+            )
+            for (member_ieee, endpoint_id) in self._zigpy_group.members.keys()
+            if member_ieee in self._zha_gateway.devices
         ]
 
-    async def async_add_members(self, member_ieee_addresses: List[EUI64]) -> None:
+    async def async_add_members(self, members: List[GroupMember]) -> None:
         """Add members to this group."""
-        if len(member_ieee_addresses) > 1:
+        if len(members) > 1:
             tasks = []
-            for ieee in member_ieee_addresses:
+            for member in members:
                 tasks.append(
-                    self._zha_gateway.devices[ieee].async_add_to_group(self.group_id)
+                    self._zha_gateway.devices[member.ieee].async_add_endpoint_to_group(
-                )
+                        member.endpoint_id, self.group_id
-            await asyncio.gather(*tasks)
-        else:
-            await self._zha_gateway.devices[
-                member_ieee_addresses[0]
-            ].async_add_to_group(self.group_id)
-
-    async def async_remove_members(self, member_ieee_addresses: List[EUI64]) -> None:
-        """Remove members from this group."""
-        if len(member_ieee_addresses) > 1:
-            tasks = []
-            for ieee in member_ieee_addresses:
-                tasks.append(
-                    self._zha_gateway.devices[ieee].async_remove_from_group(
-                        self.group_id
                     )
                 )
             await asyncio.gather(*tasks)
         else:
             await self._zha_gateway.devices[
-                member_ieee_addresses[0]
+                members[0].ieee
-            ].async_remove_from_group(self.group_id)
+            ].async_add_endpoint_to_group(members[0].endpoint_id, self.group_id)
+
+    async def async_remove_members(self, members: List[GroupMember]) -> None:
+        """Remove members from this group."""
+        if len(members) > 1:
+            tasks = []
+            for member in members:
+                tasks.append(
+                    self._zha_gateway.devices[
+                        member.ieee
+                    ].async_remove_endpoint_from_group(
+                        member.endpoint_id, self.group_id
+                    )
+                )
+            await asyncio.gather(*tasks)
+        else:
+            await self._zha_gateway.devices[
+                members[0].ieee
+            ].async_remove_endpoint_from_group(members[0].endpoint_id, self.group_id)
 
     @property
     def member_entity_ids(self) -> List[str]:
         """Return the ZHA entity ids for all entities for the members of this group."""
         all_entity_ids: List[str] = []
-        for device in self.members:
+        for member in self.members:
-            entities = async_entries_for_device(
+            entity_references = member.associated_entities
-                self._zha_gateway.ha_entity_registry, device.device_id
+            for entity_reference in entity_references:
-            )
+                all_entity_ids.append(entity_reference["entity_id"])
-            for entity in entities:
-                all_entity_ids.append(entity.entity_id)
         return all_entity_ids
 
     def get_domain_entity_ids(self, domain) -> List[str]:
         """Return entity ids from the entity domain for this group."""
         domain_entity_ids: List[str] = []
-        for device in self.members:
+        for member in self.members:
             entities = async_entries_for_device(
-                self._zha_gateway.ha_entity_registry, device.device_id
+                self._zha_gateway.ha_entity_registry, member.device.device_id
             )
             domain_entity_ids.extend(
                 [entity.entity_id for entity in entities if entity.domain == domain]
             )
         return domain_entity_ids
 
-    @callback
+    @property
-    def async_get_info(self) -> Dict[str, Any]:
+    def group_info(self) -> Dict[str, Any]:
         """Get ZHA group info."""
         group_info: Dict[str, Any] = {}
         group_info["group_id"] = self.group_id
         group_info["name"] = self.name
-        group_info["members"] = [
+        group_info["members"] = [member.member_info for member in self.members]
-            zha_device.async_get_info() for zha_device in self.members
-        ]
         return group_info
 
     def log(self, level: int, msg: str, *args):

tests/components/zha/common.py

@@ -32,7 +32,7 @@ class FakeEndpoint:
         self.model = model
         self.profile_id = zigpy.profiles.zha.PROFILE_ID
         self.device_type = None
-        self.request = AsyncMock()
+        self.request = AsyncMock(return_value=[0])
 
     def add_input_cluster(self, cluster_id):
         """Add an input cluster."""
@@ -60,6 +60,7 @@ class FakeEndpoint:
 
 
 FakeEndpoint.add_to_group = zigpy_ep.add_to_group
+FakeEndpoint.remove_from_group = zigpy_ep.remove_from_group
 
 
 def patch_cluster(cluster):

tests/components/zha/test_api.py

@@ -244,21 +244,26 @@ async def test_list_groupable_devices(zha_client, device_groupable):
     assert msg["id"] == 10
     assert msg["type"] == const.TYPE_RESULT
 
-    devices = msg["result"]
+    device_endpoints = msg["result"]
-    assert len(devices) == 1
+    assert len(device_endpoints) == 1
 
-    for device in devices:
+    for endpoint in device_endpoints:
-        assert device[ATTR_IEEE] == "01:2d:6f:00:0a:90:69:e8"
+        assert endpoint["device"][ATTR_IEEE] == "01:2d:6f:00:0a:90:69:e8"
-        assert device[ATTR_MANUFACTURER] is not None
+        assert endpoint["device"][ATTR_MANUFACTURER] is not None
-        assert device[ATTR_MODEL] is not None
+        assert endpoint["device"][ATTR_MODEL] is not None
-        assert device[ATTR_NAME] is not None
+        assert endpoint["device"][ATTR_NAME] is not None
-        assert device[ATTR_QUIRK_APPLIED] is not None
+        assert endpoint["device"][ATTR_QUIRK_APPLIED] is not None
-        assert device["entities"] is not None
+        assert endpoint["device"]["entities"] is not None
+        assert endpoint["endpoint_id"] is not None
+        assert endpoint["entities"] is not None
 
-        for entity_reference in device["entities"]:
+        for entity_reference in endpoint["device"]["entities"]:
             assert entity_reference[ATTR_NAME] is not None
             assert entity_reference["entity_id"] is not None
 
+        for entity_reference in endpoint["entities"]:
+            assert entity_reference["original_name"] is not None
+
     # Make sure there are no groupable devices when the device is unavailable
     # Make device unavailable
     device_groupable.set_available(False)
@@ -269,8 +274,8 @@ async def test_list_groupable_devices(zha_client, device_groupable):
     assert msg["id"] == 11
     assert msg["type"] == const.TYPE_RESULT
 
-    devices = msg["result"]
+    device_endpoints = msg["result"]
-    assert len(devices) == 0
+    assert len(device_endpoints) == 0
 
 
 async def test_add_group(zha_client):

tests/components/zha/test_fan.py

@@ -62,6 +62,7 @@ async def coordinator(hass, zigpy_device_mock, zha_device_joined):
         },
         ieee="00:15:8d:00:02:32:4f:32",
         nwk=0x0000,
+        node_descriptor=b"\xf8\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff",
     )
     zha_device = await zha_device_joined(zigpy_device)
     zha_device.set_available(True)

tests/components/zha/test_gateway.py

@@ -1,6 +1,8 @@
 """Test ZHA Gateway."""
+import asyncio
 import logging
 import time
+from unittest.mock import patch
 
 import pytest
 import zigpy.profiles.zha as zha
@@ -8,6 +10,7 @@ import zigpy.zcl.clusters.general as general
 import zigpy.zcl.clusters.lighting as lighting
 
 from homeassistant.components.light import DOMAIN as LIGHT_DOMAIN
+from homeassistant.components.zha.core.group import GroupMember
 
 from .common import async_enable_traffic, async_find_group_entity_id, get_zha_gateway
 
@@ -52,6 +55,7 @@ async def coordinator(hass, zigpy_device_mock, zha_device_joined):
         },
         ieee="00:15:8d:00:02:32:4f:32",
         nwk=0x0000,
+        node_descriptor=b"\xf8\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff",
     )
     zha_device = await zha_device_joined(zigpy_device)
     zha_device.set_available(True)
@@ -127,17 +131,16 @@ async def test_gateway_group_methods(hass, device_light_1, device_light_2, coord
     device_light_1._zha_gateway = zha_gateway
     device_light_2._zha_gateway = zha_gateway
     member_ieee_addresses = [device_light_1.ieee, device_light_2.ieee]
+    members = [GroupMember(device_light_1.ieee, 1), GroupMember(device_light_2.ieee, 1)]
 
     # test creating a group with 2 members
-    zha_group = await zha_gateway.async_create_zigpy_group(
+    zha_group = await zha_gateway.async_create_zigpy_group("Test Group", members)
-        "Test Group", member_ieee_addresses
-    )
     await hass.async_block_till_done()
 
     assert zha_group is not None
     assert len(zha_group.members) == 2
     for member in zha_group.members:
-        assert member.ieee in member_ieee_addresses
+        assert member.device.ieee in member_ieee_addresses
 
     entity_id = async_find_group_entity_id(hass, LIGHT_DOMAIN, zha_group)
     assert hass.states.get(entity_id) is not None
@@ -157,18 +160,24 @@ async def test_gateway_group_methods(hass, device_light_1, device_light_2, coord
 
     # test creating a group with 1 member
     zha_group = await zha_gateway.async_create_zigpy_group(
-        "Test Group", [device_light_1.ieee]
+        "Test Group", [GroupMember(device_light_1.ieee, 1)]
     )
     await hass.async_block_till_done()
 
     assert zha_group is not None
     assert len(zha_group.members) == 1
     for member in zha_group.members:
-        assert member.ieee in [device_light_1.ieee]
+        assert member.device.ieee in [device_light_1.ieee]
 
     # the group entity should not have been cleaned up
     assert entity_id not in hass.states.async_entity_ids(LIGHT_DOMAIN)
 
+    with patch("zigpy.zcl.Cluster.request", side_effect=asyncio.TimeoutError):
+        await zha_group.members[0].async_remove_from_group()
+        assert len(zha_group.members) == 1
+        for member in zha_group.members:
+            assert member.device.ieee in [device_light_1.ieee]
+
 
 async def test_updating_device_store(hass, zigpy_dev_basic, zha_dev_basic):
     """Test saving data after a delay."""

tests/components/zha/test_light.py

@@ -9,6 +9,7 @@ import zigpy.zcl.clusters.lighting as lighting
 import zigpy.zcl.foundation as zcl_f
 
 from homeassistant.components.light import DOMAIN, FLASH_LONG, FLASH_SHORT
+from homeassistant.components.zha.core.group import GroupMember
 from homeassistant.components.zha.light import FLASH_EFFECTS
 from homeassistant.const import STATE_OFF, STATE_ON, STATE_UNAVAILABLE
 import homeassistant.util.dt as dt_util
@@ -28,8 +29,8 @@ from tests.common import async_fire_time_changed
 ON = 1
 OFF = 0
 IEEE_GROUPABLE_DEVICE = "01:2d:6f:00:0a:90:69:e8"
-IEEE_GROUPABLE_DEVICE2 = "02:2d:6f:00:0a:90:69:e8"
+IEEE_GROUPABLE_DEVICE2 = "02:2d:6f:00:0a:90:69:e9"
-IEEE_GROUPABLE_DEVICE3 = "03:2d:6f:00:0a:90:69:e8"
+IEEE_GROUPABLE_DEVICE3 = "03:2d:6f:00:0a:90:69:e7"
 
 LIGHT_ON_OFF = {
     1: {
@@ -77,13 +78,14 @@ async def coordinator(hass, zigpy_device_mock, zha_device_joined):
     zigpy_device = zigpy_device_mock(
         {
             1: {
-                "in_clusters": [],
+                "in_clusters": [general.Groups.cluster_id],
                 "out_clusters": [],
                 "device_type": zha.DeviceType.COLOR_DIMMABLE_LIGHT,
             }
         },
         ieee="00:15:8d:00:02:32:4f:32",
         nwk=0x0000,
+        node_descriptor=b"\xf8\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff",
     )
     zha_device = await zha_device_joined(zigpy_device)
     zha_device.set_available(True)
@@ -109,6 +111,7 @@ async def device_light_1(hass, zigpy_device_mock, zha_device_joined):
             }
         },
         ieee=IEEE_GROUPABLE_DEVICE,
+        nwk=0xB79D,
     )
     zha_device = await zha_device_joined(zigpy_device)
     zha_device.set_available(True)
@@ -134,6 +137,7 @@ async def device_light_2(hass, zigpy_device_mock, zha_device_joined):
             }
         },
         ieee=IEEE_GROUPABLE_DEVICE2,
+        nwk=0xC79E,
     )
     zha_device = await zha_device_joined(zigpy_device)
     zha_device.set_available(True)
@@ -159,6 +163,7 @@ async def device_light_3(hass, zigpy_device_mock, zha_device_joined):
             }
         },
         ieee=IEEE_GROUPABLE_DEVICE3,
+        nwk=0xB89F,
     )
     zha_device = await zha_device_joined(zigpy_device)
     zha_device.set_available(True)
@@ -289,10 +294,12 @@ async def async_test_on_off_from_light(hass, cluster, entity_id):
     """Test on off functionality from the light."""
     # turn on at light
     await send_attributes_report(hass, cluster, {1: 0, 0: 1, 2: 3})
+    await hass.async_block_till_done()
     assert hass.states.get(entity_id).state == STATE_ON
 
     # turn off at light
     await send_attributes_report(hass, cluster, {1: 1, 0: 0, 2: 3})
+    await hass.async_block_till_done()
     assert hass.states.get(entity_id).state == STATE_OFF
 
 
@@ -300,6 +307,7 @@ async def async_test_on_from_light(hass, cluster, entity_id):
     """Test on off functionality from the light."""
     # turn on at light
     await send_attributes_report(hass, cluster, {1: -1, 0: 1, 2: 2})
+    await hass.async_block_till_done()
     assert hass.states.get(entity_id).state == STATE_ON
 
 
@@ -410,6 +418,7 @@ async def async_test_dimmer_from_light(hass, cluster, entity_id, level, expected
     await send_attributes_report(
         hass, cluster, {1: level + 10, 0: level, 2: level - 10 or 22}
     )
+    await hass.async_block_till_done()
     assert hass.states.get(entity_id).state == expected_state
     # hass uses None for brightness of 0 in state attributes
     if level == 0:
@@ -438,7 +447,23 @@ async def async_test_flash_from_hass(hass, cluster, entity_id, flash):
     )
 
 
-async def async_test_zha_group_light_entity(
+@patch(
+    "zigpy.zcl.clusters.lighting.Color.request",
+    new=AsyncMock(return_value=[sentinel.data, zcl_f.Status.SUCCESS]),
+)
+@patch(
+    "zigpy.zcl.clusters.general.Identify.request",
+    new=AsyncMock(return_value=[sentinel.data, zcl_f.Status.SUCCESS]),
+)
+@patch(
+    "zigpy.zcl.clusters.general.LevelControl.request",
+    new=AsyncMock(return_value=[sentinel.data, zcl_f.Status.SUCCESS]),
+)
+@patch(
+    "zigpy.zcl.clusters.general.OnOff.request",
+    new=AsyncMock(return_value=[sentinel.data, zcl_f.Status.SUCCESS]),
+)
+async def test_zha_group_light_entity(
     hass, device_light_1, device_light_2, device_light_3, coordinator
 ):
     """Test the light entity for a ZHA group."""
@@ -449,119 +474,180 @@ async def async_test_zha_group_light_entity(
     device_light_1._zha_gateway = zha_gateway
     device_light_2._zha_gateway = zha_gateway
     member_ieee_addresses = [device_light_1.ieee, device_light_2.ieee]
+    members = [GroupMember(device_light_1.ieee, 1), GroupMember(device_light_2.ieee, 1)]
+
+    assert coordinator.is_coordinator
 
     # test creating a group with 2 members
-    zha_group = await zha_gateway.async_create_zigpy_group(
+    zha_group = await zha_gateway.async_create_zigpy_group("Test Group", members)
-        "Test Group", member_ieee_addresses
-    )
     await hass.async_block_till_done()
 
     assert zha_group is not None
     assert len(zha_group.members) == 2
     for member in zha_group.members:
-        assert member.ieee in member_ieee_addresses
+        assert member.device.ieee in member_ieee_addresses
+        assert member.group == zha_group
+        assert member.endpoint is not None
 
-    entity_id = async_find_group_entity_id(hass, DOMAIN, zha_group)
+    device_1_entity_id = await find_entity_id(DOMAIN, device_light_1, hass)
-    assert hass.states.get(entity_id) is not None
+    device_2_entity_id = await find_entity_id(DOMAIN, device_light_2, hass)
+    device_3_entity_id = await find_entity_id(DOMAIN, device_light_3, hass)
+
+    assert (
+        device_1_entity_id != device_2_entity_id
+        and device_1_entity_id != device_3_entity_id
+    )
+    assert device_2_entity_id != device_3_entity_id
+
+    group_entity_id = async_find_group_entity_id(hass, DOMAIN, zha_group)
+    assert hass.states.get(group_entity_id) is not None
+
+    assert device_1_entity_id in zha_group.member_entity_ids
+    assert device_2_entity_id in zha_group.member_entity_ids
+    assert device_3_entity_id not in zha_group.member_entity_ids
 
     group_cluster_on_off = zha_group.endpoint[general.OnOff.cluster_id]
     group_cluster_level = zha_group.endpoint[general.LevelControl.cluster_id]
     group_cluster_identify = zha_group.endpoint[general.Identify.cluster_id]
 
-    dev1_cluster_on_off = device_light_1.endpoints[1].on_off
+    dev1_cluster_on_off = device_light_1.device.endpoints[1].on_off
-    dev2_cluster_on_off = device_light_2.endpoints[1].on_off
+    dev2_cluster_on_off = device_light_2.device.endpoints[1].on_off
-    dev3_cluster_on_off = device_light_3.endpoints[1].on_off
+    dev3_cluster_on_off = device_light_3.device.endpoints[1].on_off
+
+    dev1_cluster_level = device_light_1.device.endpoints[1].level
 
     # test that the lights were created and that they are unavailable
-    assert hass.states.get(entity_id).state == STATE_UNAVAILABLE
+    assert hass.states.get(group_entity_id).state == STATE_UNAVAILABLE
 
     # allow traffic to flow through the gateway and device
-    await async_enable_traffic(hass, zha_group.members)
+    await async_enable_traffic(hass, [device_light_1, device_light_2, device_light_3])
+    await hass.async_block_till_done()
 
     # test that the lights were created and are off
-    assert hass.states.get(entity_id).state == STATE_OFF
+    assert hass.states.get(group_entity_id).state == STATE_OFF
-
-    # test turning the lights on and off from the light
-    await async_test_on_off_from_light(hass, group_cluster_on_off, entity_id)
 
     # test turning the lights on and off from the HA
-    await async_test_on_off_from_hass(hass, group_cluster_on_off, entity_id)
+    await async_test_on_off_from_hass(hass, group_cluster_on_off, group_entity_id)
 
     # test short flashing the lights from the HA
     await async_test_flash_from_hass(
-        hass, group_cluster_identify, entity_id, FLASH_SHORT
+        hass, group_cluster_identify, group_entity_id, FLASH_SHORT
     )
 
+    # test turning the lights on and off from the light
+    await async_test_on_off_from_light(hass, dev1_cluster_on_off, group_entity_id)
+
     # test turning the lights on and off from the HA
     await async_test_level_on_off_from_hass(
-        hass, group_cluster_on_off, group_cluster_level, entity_id
+        hass, group_cluster_on_off, group_cluster_level, group_entity_id
     )
 
     # test getting a brightness change from the network
-    await async_test_on_from_light(hass, group_cluster_on_off, entity_id)
+    await async_test_on_from_light(hass, dev1_cluster_on_off, group_entity_id)
     await async_test_dimmer_from_light(
-        hass, group_cluster_level, entity_id, 150, STATE_ON
+        hass, dev1_cluster_level, group_entity_id, 150, STATE_ON
     )
 
     # test long flashing the lights from the HA
     await async_test_flash_from_hass(
-        hass, group_cluster_identify, entity_id, FLASH_LONG
+        hass, group_cluster_identify, group_entity_id, FLASH_LONG
     )
 
+    assert len(zha_group.members) == 2
     # test some of the group logic to make sure we key off states correctly
-    await dev1_cluster_on_off.on()
+    await send_attributes_report(hass, dev1_cluster_on_off, {0: 1})
-    await dev2_cluster_on_off.on()
+    await send_attributes_report(hass, dev2_cluster_on_off, {0: 1})
+    await hass.async_block_till_done()
 
     # test that group light is on
-    assert hass.states.get(entity_id).state == STATE_ON
+    assert hass.states.get(device_1_entity_id).state == STATE_ON
+    assert hass.states.get(device_2_entity_id).state == STATE_ON
+    assert hass.states.get(group_entity_id).state == STATE_ON
 
-    await dev1_cluster_on_off.off()
+    await send_attributes_report(hass, dev1_cluster_on_off, {0: 0})
+    await hass.async_block_till_done()
 
     # test that group light is still on
-    assert hass.states.get(entity_id).state == STATE_ON
+    assert hass.states.get(device_1_entity_id).state == STATE_OFF
+    assert hass.states.get(device_2_entity_id).state == STATE_ON
+    assert hass.states.get(group_entity_id).state == STATE_ON
 
-    await dev2_cluster_on_off.off()
+    await send_attributes_report(hass, dev2_cluster_on_off, {0: 0})
+    await hass.async_block_till_done()
 
     # test that group light is now off
-    assert hass.states.get(entity_id).state == STATE_OFF
+    assert hass.states.get(device_1_entity_id).state == STATE_OFF
+    assert hass.states.get(device_2_entity_id).state == STATE_OFF
+    assert hass.states.get(group_entity_id).state == STATE_OFF
 
-    await dev1_cluster_on_off.on()
+    await send_attributes_report(hass, dev1_cluster_on_off, {0: 1})
+    await hass.async_block_till_done()
 
     # test that group light is now back on
-    assert hass.states.get(entity_id).state == STATE_ON
+    assert hass.states.get(device_1_entity_id).state == STATE_ON
+    assert hass.states.get(device_2_entity_id).state == STATE_OFF
+    assert hass.states.get(group_entity_id).state == STATE_ON
 
-    # test that group light is now off
+    # turn it off to test a new member add being tracked
-    await group_cluster_on_off.off()
+    await send_attributes_report(hass, dev1_cluster_on_off, {0: 0})
-    assert hass.states.get(entity_id).state == STATE_OFF
+    await hass.async_block_till_done()
+    assert hass.states.get(device_1_entity_id).state == STATE_OFF
+    assert hass.states.get(device_2_entity_id).state == STATE_OFF
+    assert hass.states.get(group_entity_id).state == STATE_OFF
 
     # add a new member and test that his state is also tracked
-    await zha_group.async_add_members([device_light_3.ieee])
+    await zha_group.async_add_members([GroupMember(device_light_3.ieee, 1)])
-    await dev3_cluster_on_off.on()
+    await send_attributes_report(hass, dev3_cluster_on_off, {0: 1})
-    assert hass.states.get(entity_id).state == STATE_ON
+    await hass.async_block_till_done()
+    assert device_3_entity_id in zha_group.member_entity_ids
+    assert len(zha_group.members) == 3
+
+    assert hass.states.get(device_1_entity_id).state == STATE_OFF
+    assert hass.states.get(device_2_entity_id).state == STATE_OFF
+    assert hass.states.get(device_3_entity_id).state == STATE_ON
+    assert hass.states.get(group_entity_id).state == STATE_ON
 
     # make the group have only 1 member and now there should be no entity
-    await zha_group.async_remove_members([device_light_2.ieee, device_light_3.ieee])
+    await zha_group.async_remove_members(
+        [GroupMember(device_light_2.ieee, 1), GroupMember(device_light_3.ieee, 1)]
+    )
     assert len(zha_group.members) == 1
-    assert hass.states.get(entity_id).state is None
+    assert hass.states.get(group_entity_id) is None
+    assert device_2_entity_id not in zha_group.member_entity_ids
+    assert device_3_entity_id not in zha_group.member_entity_ids
+
     # make sure the entity registry entry is still there
-    assert zha_gateway.ha_entity_registry.async_get(entity_id) is not None
+    assert zha_gateway.ha_entity_registry.async_get(group_entity_id) is not None
 
     # add a member back and ensure that the group entity was created again
-    await zha_group.async_add_members([device_light_3.ieee])
+    await zha_group.async_add_members([GroupMember(device_light_3.ieee, 1)])
-    await dev3_cluster_on_off.on()
+    await send_attributes_report(hass, dev3_cluster_on_off, {0: 1})
-    assert hass.states.get(entity_id).state == STATE_ON
+    await hass.async_block_till_done()
+    assert len(zha_group.members) == 2
+    assert hass.states.get(group_entity_id).state == STATE_ON
 
     # add a 3rd member and ensure we still have an entity and we track the new one
-    await dev1_cluster_on_off.off()
+    await send_attributes_report(hass, dev1_cluster_on_off, {0: 0})
-    await dev3_cluster_on_off.off()
+    await send_attributes_report(hass, dev3_cluster_on_off, {0: 0})
-    assert hass.states.get(entity_id).state == STATE_OFF
+    await hass.async_block_till_done()
+    assert hass.states.get(group_entity_id).state == STATE_OFF
+
     # this will test that _reprobe_group is used correctly
-    await zha_group.async_add_members([device_light_2.ieee])
+    await zha_group.async_add_members(
-    await dev2_cluster_on_off.on()
+        [GroupMember(device_light_2.ieee, 1), GroupMember(coordinator.ieee, 1)]
-    assert hass.states.get(entity_id).state == STATE_ON
+    )
+    await send_attributes_report(hass, dev2_cluster_on_off, {0: 1})
+    await hass.async_block_till_done()
+    assert len(zha_group.members) == 4
+    assert hass.states.get(group_entity_id).state == STATE_ON
+
+    await zha_group.async_remove_members([GroupMember(coordinator.ieee, 1)])
+    await hass.async_block_till_done()
+    assert hass.states.get(group_entity_id).state == STATE_ON
+    assert len(zha_group.members) == 3
 
     # remove the group and ensure that there is no entity and that the entity registry is cleaned up
-    assert zha_gateway.ha_entity_registry.async_get(entity_id) is not None
+    assert zha_gateway.ha_entity_registry.async_get(group_entity_id) is not None
     await zha_gateway.async_remove_zigpy_group(zha_group.group_id)
-    assert hass.states.get(entity_id).state is None
+    assert hass.states.get(group_entity_id) is None
-    assert zha_gateway.ha_entity_registry.async_get(entity_id) is None
+    assert zha_gateway.ha_entity_registry.async_get(group_entity_id) is None

tests/components/zha/test_switch.py

@@ -53,6 +53,7 @@ async def coordinator(hass, zigpy_device_mock, zha_device_joined):
         },
         ieee="00:15:8d:00:02:32:4f:32",
         nwk=0x0000,
+        node_descriptor=b"\xf8\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff",
     )
     zha_device = await zha_device_joined(zigpy_device)
     zha_device.set_available(True)