Add group entity support to ZHA

homeassistant/components/zha/__init__.py

@@ -130,7 +130,7 @@ async def async_zha_shutdown(event):
         await zha_data[DATA_ZHA_GATEWAY].async_update_device_storage()
 
     hass.bus.async_listen_once(ha_const.EVENT_HOMEASSISTANT_STOP, async_zha_shutdown)
-    asyncio.create_task(async_load_entities(hass, config_entry))
+    asyncio.create_task(async_load_entities(hass))
     return True
 
 
@@ -150,11 +150,9 @@ async def async_unload_entry(hass, config_entry):
     return True
 
 
-async def async_load_entities(
-    hass: HomeAssistantType, config_entry: config_entries.ConfigEntry
-) -> None:
+async def async_load_entities(hass: HomeAssistantType) -> None:
     """Load entities after integration was setup."""
-    await hass.data[DATA_ZHA][DATA_ZHA_GATEWAY].async_prepare_entities()
+    await hass.data[DATA_ZHA][DATA_ZHA_GATEWAY].async_initialize_devices_and_entities()
     to_setup = hass.data[DATA_ZHA][DATA_ZHA_PLATFORM_LOADED]
     results = await asyncio.gather(*to_setup, return_exceptions=True)
     for res in results:

homeassistant/components/zha/core/const.py

@@ -23,6 +23,7 @@
 ATTR_DEVICE_IEEE = "device_ieee"
 ATTR_DEVICE_TYPE = "device_type"
 ATTR_ENDPOINT_ID = "endpoint_id"
+ATTR_ENTITY_DOMAIN = "entity_domain"
 ATTR_IEEE = "ieee"
 ATTR_LAST_SEEN = "last_seen"
 ATTR_LEVEL = "level"
@@ -207,6 +208,7 @@ def list(cls):
 SIGNAL_SET_LEVEL = "set_level"
 SIGNAL_STATE_ATTR = "update_state_attribute"
 SIGNAL_UPDATE_DEVICE = "{}_zha_update_device"
+SIGNAL_REMOVE_GROUP = "remove_group"
 
 UNKNOWN = "unknown"
 UNKNOWN_MANUFACTURER = "unk_manufacturer"

homeassistant/components/zha/core/device.py

@@ -373,7 +373,7 @@ async def async_configure(self):
         self.debug("started configuration")
         await self._channels.async_configure()
         self.debug("completed configuration")
-        entry = self.gateway.zha_storage.async_create_or_update(self)
+        entry = self.gateway.zha_storage.async_create_or_update_device(self)
         self.debug("stored in registry: %s", entry)
 
         if self._channels.identify_ch is not None:

homeassistant/components/zha/core/discovery.py

@@ -1,10 +1,12 @@
 """Device discovery functions for Zigbee Home Automation."""
 
+from collections import Counter
 import logging
 from typing import Callable, List, Tuple
 
 from homeassistant import const as ha_const
 from homeassistant.core import callback
+from homeassistant.helpers.entity_registry import async_entries_for_device
 from homeassistant.helpers.typing import HomeAssistantType
 
 from . import const as zha_const, registries as zha_regs, typing as zha_typing
@@ -157,4 +159,102 @@ def initialize(self, hass: HomeAssistantType) -> None:
             self._device_configs.update(overrides)
 
 
+class GroupProbe:
+    """Determine the appropriate component for a group."""
+
+    def __init__(self):
+        """Initialize instance."""
+        self._hass = None
+
+    def initialize(self, hass: HomeAssistantType) -> None:
+        """Initialize the group probe."""
+        self._hass = hass
+
+    @callback
+    def discover_group_entities(self, group: zha_typing.ZhaGroupType) -> None:
+        """Process a group and create any entities that are needed."""
+        # only create a group entity if there are 2 or more members in a group
+        if len(group.members) < 2:
+            _LOGGER.debug(
+                "Group: %s:0x%04x has less than 2 members - skipping entity discovery",
+                group.name,
+                group.group_id,
+            )
+            return
+
+        if group.entity_domain is None:
+            _LOGGER.debug(
+                "Group: %s:0x%04x has no user set entity domain - attempting entity domain discovery",
+                group.name,
+                group.group_id,
+            )
+            group.entity_domain = GroupProbe.determine_default_entity_domain(
+                self._hass, group
+            )
+
+        if group.entity_domain is None:
+            return
+
+        _LOGGER.debug(
+            "Group: %s:0x%04x has an entity domain of: %s after discovery",
+            group.name,
+            group.group_id,
+            group.entity_domain,
+        )
+
+        zha_gateway = self._hass.data[zha_const.DATA_ZHA][zha_const.DATA_ZHA_GATEWAY]
+        entity_class = zha_regs.ZHA_ENTITIES.get_group_entity(group.entity_domain)
+        if entity_class is None:
+            return
+
+        self._hass.data[zha_const.DATA_ZHA][group.entity_domain].append(
+            (
+                entity_class,
+                (
+                    group.domain_entity_ids,
+                    f"{group.entity_domain}_group_{group.group_id}",
+                    group.group_id,
+                    zha_gateway.coordinator_zha_device,
+                ),
+            )
+        )
+
+    @staticmethod
+    def determine_default_entity_domain(
+        hass: HomeAssistantType, group: zha_typing.ZhaGroupType
+    ):
+        """Determine the default entity domain for this group."""
+        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 None
+
+        zha_gateway = hass.data[zha_const.DATA_ZHA][zha_const.DATA_ZHA_GATEWAY]
+        all_domain_occurrences = []
+        for device in group.members:
+            entities = async_entries_for_device(
+                zha_gateway.ha_entity_registry, device.device_id
+            )
+            all_domain_occurrences.extend(
+                [
+                    entity.domain
+                    for entity in entities
+                    if entity.domain in zha_regs.GROUP_ENTITY_DOMAINS
+                ]
+            )
+        counts = Counter(all_domain_occurrences)
+        domain = counts.most_common(1)[0][0]
+        _LOGGER.debug(
+            "The default entity domain is: %s for group: %s:0x%04x",
+            domain,
+            group.name,
+            group.group_id,
+        )
+        return domain
+
+
 PROBE = ProbeEndpoint()
+GROUP_PROBE = GroupProbe()