Add identify buttons to ZHA devices (#61495)

* Identify buttons

* clean up and add test

* use Platform

* update device list

* Only 1 identify button per device

* cleanup press until the need arises for the branch

* make imports relative

tests/components/zha/test_button.py

@@ -0,0 +1,89 @@
+"""Test ZHA button."""
+from unittest.mock import patch
+
+from freezegun import freeze_time
+import pytest
+from zigpy.const import SIG_EP_PROFILE
+import zigpy.profiles.zha as zha
+import zigpy.zcl.clusters.general as general
+import zigpy.zcl.clusters.security as security
+import zigpy.zcl.foundation as zcl_f
+
+from homeassistant.components.button import DOMAIN, ButtonDeviceClass
+from homeassistant.components.button.const import SERVICE_PRESS
+from homeassistant.const import (
+    ATTR_DEVICE_CLASS,
+    ATTR_ENTITY_ID,
+    ENTITY_CATEGORY_DIAGNOSTIC,
+    STATE_UNKNOWN,
+)
+from homeassistant.helpers import entity_registry as er
+
+from .common import find_entity_id
+from .conftest import SIG_EP_INPUT, SIG_EP_OUTPUT, SIG_EP_TYPE
+
+from tests.common import mock_coro
+
+
+@pytest.fixture
+async def contact_sensor(hass, zigpy_device_mock, zha_device_joined_restored):
+    """Contact sensor fixture."""
+
+    zigpy_device = zigpy_device_mock(
+        {
+            1: {
+                SIG_EP_INPUT: [
+                    general.Basic.cluster_id,
+                    general.Identify.cluster_id,
+                    security.IasZone.cluster_id,
+                ],
+                SIG_EP_OUTPUT: [],
+                SIG_EP_TYPE: zha.DeviceType.IAS_ZONE,
+                SIG_EP_PROFILE: zha.PROFILE_ID,
+            }
+        },
+    )
+
+    zha_device = await zha_device_joined_restored(zigpy_device)
+    return zha_device, zigpy_device.endpoints[1].identify
+
+
+@freeze_time("2021-11-04 17:37:00", tz_offset=-1)
+async def test_button(hass, contact_sensor):
+    """Test zha button platform."""
+
+    entity_registry = er.async_get(hass)
+    zha_device, cluster = contact_sensor
+    assert cluster is not None
+    entity_id = await find_entity_id(DOMAIN, zha_device, hass)
+    assert entity_id is not None
+
+    state = hass.states.get(entity_id)
+    assert state
+    assert state.state == STATE_UNKNOWN
+    assert state.attributes[ATTR_DEVICE_CLASS] == ButtonDeviceClass.UPDATE
+
+    entry = entity_registry.async_get(entity_id)
+    assert entry
+    assert entry.entity_category == ENTITY_CATEGORY_DIAGNOSTIC
+
+    with patch(
+        "zigpy.zcl.Cluster.request",
+        return_value=mock_coro([0x00, zcl_f.Status.SUCCESS]),
+    ):
+        await hass.services.async_call(
+            DOMAIN,
+            SERVICE_PRESS,
+            {ATTR_ENTITY_ID: entity_id},
+            blocking=True,
+        )
+        await hass.async_block_till_done()
+        assert len(cluster.request.mock_calls) == 1
+        assert cluster.request.call_args[0][0] is False
+        assert cluster.request.call_args[0][1] == 0
+        assert cluster.request.call_args[0][3] == 5  # duration in seconds
+
+    state = hass.states.get(entity_id)
+    assert state
+    assert state.state == "2021-11-04T16:37:00+00:00"
+    assert state.attributes[ATTR_DEVICE_CLASS] == ButtonDeviceClass.UPDATE

tests/components/zha/test_discover.py

@@ -125,17 +125,25 @@ async def test_devices(
         ch.id for pool in zha_dev.channels.pools for ch in pool.client_channels.values()
     }
     assert event_channels == set(device[DEV_SIG_EVT_CHANNELS])
-
+    # we need to probe the class create entity factory so we need to reset this to get accurate results
+    zha_regs.ZHA_ENTITIES.clean_up()
     # build a dict of entity_class -> (component, unique_id, channels) tuple
     ha_ent_info = {}
+    created_entity_count = 0
     for call in _dispatch.call_args_list:
         _, component, entity_cls, unique_id, channels = call[0]
-        unique_id_head = UNIQUE_ID_HD.match(unique_id).group(0)  # ieee + endpoint_id
-        ha_ent_info[(unique_id_head, entity_cls.__name__)] = (
-            component,
-            unique_id,
-            channels,
-        )
+        # the factory can return None. We filter these out to get an accurate created entity count
+        response = entity_cls.create_entity(unique_id, zha_dev, channels)
+        if response:
+            created_entity_count += 1
+            unique_id_head = UNIQUE_ID_HD.match(unique_id).group(
+                0
+            )  # ieee + endpoint_id
+            ha_ent_info[(unique_id_head, entity_cls.__name__)] = (
+                component,
+                unique_id,
+                channels,
+            )
 
     for comp_id, ent_info in device[DEV_SIG_ENT_MAP].items():
         component, unique_id = comp_id
@@ -156,7 +164,7 @@ async def test_devices(
         assert unique_id.startswith(ha_unique_id)
         assert {ch.name for ch in ha_channels} == set(ent_info[DEV_SIG_CHANNELS])
 
-    assert _dispatch.call_count == len(device[DEV_SIG_ENT_MAP])
+    assert created_entity_count == len(device[DEV_SIG_ENT_MAP])
 
     entity_ids = hass_disable_services.states.async_entity_ids()
     await hass_disable_services.async_block_till_done()
@@ -298,7 +306,6 @@ async def test_discover_endpoint(device_info, channels_mock, hass):
     assert device_info[DEV_SIG_EVT_CHANNELS] == sorted(
         ch.id for pool in channels.pools for ch in pool.client_channels.values()
     )
-    assert new_ent.call_count == len(list(device_info[DEV_SIG_ENT_MAP].values()))
 
     # build a dict of entity_class -> (component, unique_id, channels) tuple
     ha_ent_info = {}
@@ -326,8 +333,6 @@ async def test_discover_endpoint(device_info, channels_mock, hass):
         assert unique_id.startswith(ha_unique_id)
         assert {ch.name for ch in ha_channels} == set(ent_info[DEV_SIG_CHANNELS])
 
-    assert new_ent.call_count == len(device_info[DEV_SIG_ENT_MAP])
-
 
 def _ch_mock(cluster):
     """Return mock of a channel with a cluster."""