Add ZHA group API (#29641)

* add skeleton to retrieve zigbee groups
* get single group
* add a group
* return group members with group
* add comment
* fix group members
* add function to add device to group
* add group members
* add remove from group method
* add api to remove members from group
* add remove groups method
* clean up group add and remove
* fix remove group
* fix remove groups
* add api to get only groupable devices
* change var init
* add tests
* address review comment

homeassistant/components/zha/api.py

@@ -23,6 +23,7 @@ from .core.const import (
     ATTR_ENDPOINT_ID,
     ATTR_LEVEL,
     ATTR_MANUFACTURER,
+    ATTR_MEMBERS,
     ATTR_NAME,
     ATTR_VALUE,
     ATTR_WARNING_DEVICE_DURATION,
@@ -39,6 +40,9 @@ from .core.const import (
     DATA_ZHA,
     DATA_ZHA_GATEWAY,
     DOMAIN,
+    GROUP_ID,
+    GROUP_IDS,
+    GROUP_NAME,
     MFG_CLUSTER_ID_START,
     WARNING_DEVICE_MODE_EMERGENCY,
     WARNING_DEVICE_SOUND_HIGH,
@@ -211,6 +215,34 @@ async def websocket_get_devices(hass, connection, msg):
     connection.send_result(msg[ID], devices)
 
 
+@websocket_api.require_admin
+@websocket_api.async_response
+@websocket_api.websocket_command({vol.Required(TYPE): "zha/devices/groupable"})
+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]
+    ha_device_registry = await async_get_registry(hass)
+
+    devices = []
+    for device in zha_gateway.devices.values():
+        if device.is_groupable:
+            devices.append(
+                async_get_device_info(
+                    hass, device, ha_device_registry=ha_device_registry
+                )
+            )
+    connection.send_result(msg[ID], devices)
+
+
+@websocket_api.require_admin
+@websocket_api.async_response
+@websocket_api.websocket_command({vol.Required(TYPE): "zha/groups"})
+async def websocket_get_groups(hass, connection, msg):
+    """Get ZHA groups."""
+    groups = await get_groups(hass)
+    connection.send_result(msg[ID], groups)
+
+
 @websocket_api.require_admin
 @websocket_api.async_response
 @websocket_api.websocket_command(
@@ -236,6 +268,161 @@ async def websocket_get_device(hass, connection, msg):
     connection.send_result(msg[ID], device)
 
 
+@websocket_api.require_admin
+@websocket_api.async_response
+@websocket_api.websocket_command(
+    {vol.Required(TYPE): "zha/group", vol.Required(GROUP_ID): cv.positive_int}
+)
+async def websocket_get_group(hass, connection, msg):
+    """Get ZHA group."""
+    zha_gateway = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY]
+    ha_device_registry = await async_get_registry(hass)
+    group_id = msg[GROUP_ID]
+    group = None
+
+    if group_id in zha_gateway.application_controller.groups:
+        group = async_get_group_info(
+            hass,
+            zha_gateway,
+            zha_gateway.application_controller.groups[group_id],
+            ha_device_registry,
+        )
+    if not group:
+        connection.send_message(
+            websocket_api.error_message(
+                msg[ID], websocket_api.const.ERR_NOT_FOUND, "ZHA Group not found"
+            )
+        )
+        return
+    connection.send_result(msg[ID], group)
+
+
+@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]),
+    }
+)
+async def websocket_add_group(hass, connection, msg):
+    """Add a new ZHA group."""
+    zha_gateway = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY]
+    ha_device_registry = await async_get_registry(hass)
+    group_id = len(zha_gateway.application_controller.groups) + 1
+    group_name = msg[GROUP_NAME]
+    zigpy_group = async_get_group_by_name(zha_gateway, group_name)
+    ret_group = None
+    members = msg.get(ATTR_MEMBERS)
+
+    # guard against group already existing
+    if zigpy_group is None:
+        zigpy_group = zha_gateway.application_controller.groups.add_group(
+            group_id, group_name
+        )
+        if members is not None:
+            tasks = []
+            for ieee in members:
+                tasks.append(zha_gateway.devices[ieee].async_add_to_group(group_id))
+            await asyncio.gather(*tasks)
+    ret_group = async_get_group_info(hass, zha_gateway, zigpy_group, ha_device_registry)
+    connection.send_result(msg[ID], ret_group)
+
+
+@websocket_api.require_admin
+@websocket_api.async_response
+@websocket_api.websocket_command(
+    {
+        vol.Required(TYPE): "zha/group/remove",
+        vol.Required(GROUP_IDS): vol.All(cv.ensure_list, [cv.positive_int]),
+    }
+)
+async def websocket_remove_groups(hass, connection, msg):
+    """Remove the specified ZHA groups."""
+    zha_gateway = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY]
+    groups = zha_gateway.application_controller.groups
+    group_ids = msg[GROUP_IDS]
+
+    if len(group_ids) > 1:
+        tasks = []
+        for group_id in group_ids:
+            tasks.append(remove_group(groups[group_id], zha_gateway))
+        await asyncio.gather(*tasks)
+    else:
+        await remove_group(groups[group_ids[0]], zha_gateway)
+    ret_groups = await get_groups(hass)
+    connection.send_result(msg[ID], ret_groups)
+
+
+@websocket_api.require_admin
+@websocket_api.async_response
+@websocket_api.websocket_command(
+    {
+        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]),
+    }
+)
+async def websocket_add_group_members(hass, connection, msg):
+    """Add members to a ZHA group."""
+    zha_gateway = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY]
+    ha_device_registry = await async_get_registry(hass)
+    group_id = msg[GROUP_ID]
+    members = msg[ATTR_MEMBERS]
+    zigpy_group = None
+
+    if group_id in zha_gateway.application_controller.groups:
+        zigpy_group = zha_gateway.application_controller.groups[group_id]
+        tasks = []
+        for ieee in members:
+            tasks.append(zha_gateway.devices[ieee].async_add_to_group(group_id))
+        await asyncio.gather(*tasks)
+    if not zigpy_group:
+        connection.send_message(
+            websocket_api.error_message(
+                msg[ID], websocket_api.const.ERR_NOT_FOUND, "ZHA Group not found"
+            )
+        )
+        return
+    ret_group = async_get_group_info(hass, zha_gateway, zigpy_group, ha_device_registry)
+    connection.send_result(msg[ID], ret_group)
+
+
+@websocket_api.require_admin
+@websocket_api.async_response
+@websocket_api.websocket_command(
+    {
+        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]),
+    }
+)
+async def websocket_remove_group_members(hass, connection, msg):
+    """Remove members from a ZHA group."""
+    zha_gateway = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY]
+    ha_device_registry = await async_get_registry(hass)
+    group_id = msg[GROUP_ID]
+    members = msg[ATTR_MEMBERS]
+    zigpy_group = None
+
+    if group_id in zha_gateway.application_controller.groups:
+        zigpy_group = zha_gateway.application_controller.groups[group_id]
+        tasks = []
+        for ieee in members:
+            tasks.append(zha_gateway.devices[ieee].async_remove_from_group(group_id))
+        await asyncio.gather(*tasks)
+    if not zigpy_group:
+        connection.send_message(
+            websocket_api.error_message(
+                msg[ID], websocket_api.const.ERR_NOT_FOUND, "ZHA Group not found"
+            )
+        )
+        return
+    ret_group = async_get_group_info(hass, zha_gateway, zigpy_group, ha_device_registry)
+    connection.send_result(msg[ID], ret_group)
+
+
 @callback
 def async_get_device_info(hass, device, ha_device_registry=None):
     """Get ZHA device."""
@@ -261,6 +448,62 @@ def async_get_device_info(hass, device, ha_device_registry=None):
     return ret_device
 
 
+async def get_groups(hass,):
+    """Get ZHA Groups."""
+    zha_gateway = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY]
+    ha_device_registry = await async_get_registry(hass)
+
+    groups = []
+    for group in zha_gateway.application_controller.groups.values():
+        groups.append(
+            async_get_group_info(hass, zha_gateway, group, ha_device_registry)
+        )
+    return groups
+
+
+async def remove_group(group, zha_gateway):
+    """Remove ZHA Group."""
+    if group.members:
+        tasks = []
+        for member_ieee in group.members.keys():
+            if member_ieee[0] in zha_gateway.devices:
+                tasks.append(
+                    zha_gateway.devices[member_ieee[0]].async_remove_from_group(
+                        group.group_id
+                    )
+                )
+        await asyncio.gather(*tasks)
+    else:
+        zha_gateway.application_controller.groups.pop(group.group_id)
+
+
+@callback
+def async_get_group_info(hass, zha_gateway, group, ha_device_registry):
+    """Get ZHA group."""
+    ret_group = {}
+    ret_group["group_id"] = group.group_id
+    ret_group["name"] = group.name
+    ret_group["members"] = [
+        async_get_device_info(
+            hass,
+            zha_gateway.get_device(member_ieee[0]),
+            ha_device_registry=ha_device_registry,
+        )
+        for member_ieee in group.members.keys()
+        if member_ieee[0] in zha_gateway.devices
+    ]
+    return ret_group
+
+
+@callback
+def async_get_group_by_name(zha_gateway, group_name):
+    """Get ZHA group by name."""
+    for group in zha_gateway.application_controller.groups.values():
+        if group.name == group_name:
+            return group
+    return None
+
+
 @websocket_api.require_admin
 @websocket_api.async_response
 @websocket_api.websocket_command(
@@ -785,7 +1028,14 @@ def async_load_api(hass):
 
     websocket_api.async_register_command(hass, websocket_permit_devices)
     websocket_api.async_register_command(hass, websocket_get_devices)
+    websocket_api.async_register_command(hass, websocket_get_groupable_devices)
+    websocket_api.async_register_command(hass, websocket_get_groups)
     websocket_api.async_register_command(hass, websocket_get_device)
+    websocket_api.async_register_command(hass, websocket_get_group)
+    websocket_api.async_register_command(hass, websocket_add_group)
+    websocket_api.async_register_command(hass, websocket_remove_groups)
+    websocket_api.async_register_command(hass, websocket_add_group_members)
+    websocket_api.async_register_command(hass, websocket_remove_group_members)
     websocket_api.async_register_command(hass, websocket_reconfigure_node)
     websocket_api.async_register_command(hass, websocket_device_clusters)
     websocket_api.async_register_command(hass, websocket_device_cluster_attributes)

homeassistant/components/zha/core/const.py

@@ -24,6 +24,7 @@ ATTR_LEVEL = "level"
 ATTR_LQI = "lqi"
 ATTR_MANUFACTURER = "manufacturer"
 ATTR_MANUFACTURER_CODE = "manufacturer_code"
+ATTR_MEMBERS = "members"
 ATTR_MODEL = "model"
 ATTR_NAME = "name"
 ATTR_NWK = "nwk"
@@ -105,6 +106,10 @@ DISCOVERY_KEY = "zha_discovery_info"
 
 DOMAIN = "zha"
 
+GROUP_ID = "group_id"
+GROUP_IDS = "group_ids"
+GROUP_NAME = "group_name"
+
 MFG_CLUSTER_ID_START = 0xFC00
 
 POWER_MAINS_POWERED = "Mains"

homeassistant/components/zha/core/device.py

@@ -13,6 +13,7 @@ import time
 import zigpy.exceptions
 from zigpy.profiles import zha, zll
 import zigpy.quirks
+from zigpy.zcl.clusters.general import Groups
 
 from homeassistant.core import callback
 from homeassistant.helpers.dispatcher import (
@@ -179,6 +180,17 @@ class ZHADevice(LogMixin):
         """Return true if this device is an end device."""
         return self._zigpy_device.node_desc.is_end_device
 
+    @property
+    def is_groupable(self):
+        """Return true if this device has a group cluster."""
+        if not self.available:
+            return False
+        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 gateway(self):
         """Return the gateway for this device."""
@@ -506,6 +518,14 @@ class ZHADevice(LogMixin):
         )
         return response
 
+    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)
+
+    async def async_remove_from_group(self, group_id):
+        """Remove this device from the provided zigbee group."""
+        await self._zigpy_device.remove_from_group(group_id)
+
     def log(self, level, msg, *args):
         """Log a message."""
         msg = "[%s](%s): " + msg

tests/components/zha/common.py

@@ -93,6 +93,8 @@ class FakeDevice:
         self.manufacturer = manufacturer
         self.model = model
         self.node_desc = zigpy.zdo.types.NodeDescriptor()
+        self.add_to_group = CoroutineMock()
+        self.remove_from_group = CoroutineMock()
 
 
 def make_device(

tests/components/zha/conftest.py

@@ -1,7 +1,10 @@
 """Test configuration for the ZHA component."""
+from unittest import mock
 from unittest.mock import patch
 
 import pytest
+import zigpy
+from zigpy.application import ControllerApplication
 
 from homeassistant import config_entries
 from homeassistant.components.zha.core.const import COMPONENTS, DATA_ZHA, DOMAIN
@@ -12,6 +15,9 @@ from homeassistant.helpers.device_registry import async_get_registry as get_dev_
 
 from .common import async_setup_entry
 
+FIXTURE_GRP_ID = 0x1001
+FIXTURE_GRP_NAME = "fixture group"
+
 
 @pytest.fixture(name="config_entry")
 def config_entry_fixture(hass):
@@ -43,6 +49,11 @@ async def zha_gateway_fixture(hass, config_entry):
     gateway = ZHAGateway(hass, {}, config_entry)
     gateway.zha_storage = zha_storage
     gateway.ha_device_registry = dev_reg
+    gateway.application_controller = mock.MagicMock(spec_set=ControllerApplication)
+    groups = zigpy.group.Groups(gateway.application_controller)
+    groups.listener_event = mock.MagicMock()
+    groups.add_group(FIXTURE_GRP_ID, FIXTURE_GRP_NAME, suppress_event=True)
+    gateway.application_controller.groups = groups
     return gateway
 
 

tests/components/zha/test_api.py

@@ -1,7 +1,10 @@
 """Test ZHA API."""
+
 import pytest
+import zigpy
 import zigpy.zcl.clusters.general as general
 
+from homeassistant.components.light import DOMAIN as light_domain
 from homeassistant.components.switch import DOMAIN
 from homeassistant.components.websocket_api import const
 from homeassistant.components.zha.api import ID, TYPE, async_load_api
@@ -15,9 +18,13 @@ from homeassistant.components.zha.core.const import (
     ATTR_NAME,
     ATTR_QUIRK_APPLIED,
     CLUSTER_TYPE_IN,
+    GROUP_ID,
+    GROUP_IDS,
+    GROUP_NAME,
 )
 
 from .common import async_init_zigpy_device
+from .conftest import FIXTURE_GRP_ID, FIXTURE_GRP_NAME
 
 
 @pytest.fixture
@@ -36,9 +43,22 @@ async def zha_client(hass, config_entry, zha_gateway, hass_ws_client):
         zha_gateway,
     )
 
+    await async_init_zigpy_device(
+        hass,
+        [general.OnOff.cluster_id, general.Basic.cluster_id, general.Groups.cluster_id],
+        [],
+        zigpy.profiles.zha.DeviceType.ON_OFF_LIGHT,
+        zha_gateway,
+        manufacturer="FakeGroupManufacturer",
+        model="FakeGroupModel",
+        ieee="01:2d:6f:00:0a:90:69:e8",
+    )
+
     # load up switch domain
     await hass.config_entries.async_forward_entry_setup(config_entry, DOMAIN)
     await hass.async_block_till_done()
+    await hass.config_entries.async_forward_entry_setup(config_entry, light_domain)
+    await hass.async_block_till_done()
 
     return await hass_ws_client(hass)
 
@@ -114,15 +134,17 @@ async def test_device_cluster_commands(hass, config_entry, zha_gateway, zha_clie
 
 
 async def test_list_devices(hass, config_entry, zha_gateway, zha_client):
-    """Test getting entity cluster commands."""
+    """Test getting zha devices."""
     await zha_client.send_json({ID: 5, TYPE: "zha/devices"})
 
     msg = await zha_client.receive_json()
 
     devices = msg["result"]
-    assert len(devices) == 1
+    assert len(devices) == 2
 
+    msg_id = 100
     for device in devices:
+        msg_id += 1
         assert device[ATTR_IEEE] is not None
         assert device[ATTR_MANUFACTURER] is not None
         assert device[ATTR_MODEL] is not None
@@ -135,7 +157,7 @@ async def test_list_devices(hass, config_entry, zha_gateway, zha_client):
             assert entity_reference["entity_id"] is not None
 
         await zha_client.send_json(
-            {ID: 6, TYPE: "zha/device", ATTR_IEEE: device[ATTR_IEEE]}
+            {ID: msg_id, TYPE: "zha/device", ATTR_IEEE: device[ATTR_IEEE]}
         )
         msg = await zha_client.receive_json()
         device2 = msg["result"]
@@ -152,3 +174,151 @@ async def test_device_not_found(hass, config_entry, zha_gateway, zha_client):
     assert msg["type"] == const.TYPE_RESULT
     assert not msg["success"]
     assert msg["error"]["code"] == const.ERR_NOT_FOUND
+
+
+async def test_list_groups(hass, config_entry, zha_gateway, zha_client):
+    """Test getting zha zigbee groups."""
+    await zha_client.send_json({ID: 7, TYPE: "zha/groups"})
+
+    msg = await zha_client.receive_json()
+    assert msg["id"] == 7
+    assert msg["type"] == const.TYPE_RESULT
+
+    groups = msg["result"]
+    assert len(groups) == 1
+
+    for group in groups:
+        assert group["group_id"] == FIXTURE_GRP_ID
+        assert group["name"] == FIXTURE_GRP_NAME
+        assert group["members"] == []
+
+
+async def test_get_group(hass, config_entry, zha_gateway, zha_client):
+    """Test getting a specific zha zigbee group."""
+    await zha_client.send_json({ID: 8, TYPE: "zha/group", GROUP_ID: FIXTURE_GRP_ID})
+
+    msg = await zha_client.receive_json()
+    assert msg["id"] == 8
+    assert msg["type"] == const.TYPE_RESULT
+
+    group = msg["result"]
+    assert group is not None
+    assert group["group_id"] == FIXTURE_GRP_ID
+    assert group["name"] == FIXTURE_GRP_NAME
+    assert group["members"] == []
+
+
+async def test_get_group_not_found(hass, config_entry, zha_gateway, zha_client):
+    """Test not found response from get group API."""
+    await zha_client.send_json({ID: 9, TYPE: "zha/group", GROUP_ID: 1234567})
+
+    msg = await zha_client.receive_json()
+
+    assert msg["id"] == 9
+    assert msg["type"] == const.TYPE_RESULT
+    assert not msg["success"]
+    assert msg["error"]["code"] == const.ERR_NOT_FOUND
+
+
+async def test_list_groupable_devices(hass, config_entry, zha_gateway, zha_client):
+    """Test getting zha devices that have a group cluster."""
+
+    # Make device available
+    zha_gateway.devices[
+        zigpy.types.EUI64.convert("01:2d:6f:00:0a:90:69:e8")
+    ].set_available(True)
+
+    await zha_client.send_json({ID: 10, TYPE: "zha/devices/groupable"})
+
+    msg = await zha_client.receive_json()
+    assert msg["id"] == 10
+    assert msg["type"] == const.TYPE_RESULT
+
+    devices = msg["result"]
+    assert len(devices) == 1
+
+    for device in devices:
+        assert device[ATTR_IEEE] == "01:2d:6f:00:0a:90:69:e8"
+        assert device[ATTR_MANUFACTURER] is not None
+        assert device[ATTR_MODEL] is not None
+        assert device[ATTR_NAME] is not None
+        assert device[ATTR_QUIRK_APPLIED] is not None
+        assert device["entities"] is not None
+
+        for entity_reference in device["entities"]:
+            assert entity_reference[ATTR_NAME] is not None
+            assert entity_reference["entity_id"] is not None
+
+    # Make sure there are no groupable devices when the device is unavailable
+    # Make device unavailable
+    zha_gateway.devices[
+        zigpy.types.EUI64.convert("01:2d:6f:00:0a:90:69:e8")
+    ].set_available(False)
+
+    await zha_client.send_json({ID: 11, TYPE: "zha/devices/groupable"})
+
+    msg = await zha_client.receive_json()
+    assert msg["id"] == 11
+    assert msg["type"] == const.TYPE_RESULT
+
+    devices = msg["result"]
+    assert len(devices) == 0
+
+
+async def test_add_group(hass, config_entry, zha_gateway, zha_client):
+    """Test adding and getting a new zha zigbee group."""
+    await zha_client.send_json({ID: 12, TYPE: "zha/group/add", GROUP_NAME: "new_group"})
+
+    msg = await zha_client.receive_json()
+    assert msg["id"] == 12
+    assert msg["type"] == const.TYPE_RESULT
+
+    added_group = msg["result"]
+
+    assert added_group["name"] == "new_group"
+    assert added_group["members"] == []
+
+    await zha_client.send_json({ID: 13, TYPE: "zha/groups"})
+
+    msg = await zha_client.receive_json()
+    assert msg["id"] == 13
+    assert msg["type"] == const.TYPE_RESULT
+
+    groups = msg["result"]
+    assert len(groups) == 2
+
+    for group in groups:
+        assert group["name"] == FIXTURE_GRP_NAME or group["name"] == "new_group"
+
+
+async def test_remove_group(hass, config_entry, zha_gateway, zha_client):
+    """Test removing a new zha zigbee group."""
+
+    await zha_client.send_json({ID: 14, TYPE: "zha/groups"})
+
+    msg = await zha_client.receive_json()
+    assert msg["id"] == 14
+    assert msg["type"] == const.TYPE_RESULT
+
+    groups = msg["result"]
+    assert len(groups) == 1
+
+    await zha_client.send_json(
+        {ID: 15, TYPE: "zha/group/remove", GROUP_IDS: [FIXTURE_GRP_ID]}
+    )
+
+    msg = await zha_client.receive_json()
+    assert msg["id"] == 15
+    assert msg["type"] == const.TYPE_RESULT
+
+    groups_remaining = msg["result"]
+    assert len(groups_remaining) == 0
+
+    await zha_client.send_json({ID: 16, TYPE: "zha/groups"})
+
+    msg = await zha_client.receive_json()
+    assert msg["id"] == 16
+    assert msg["type"] == const.TYPE_RESULT
+
+    groups = msg["result"]
+    assert len(groups) == 0