Add device counter entities to ZHA (#111175)

* Add counter entities to the ZHA coordinator device

* rework to prepare for non coordinator device counters

* counter entity test

* update log lines

* disable by default

homeassistant/components/zha/core/device.py

@@ -33,7 +33,7 @@ from homeassistant.helpers.dispatcher import (
 )
 from homeassistant.helpers.event import async_track_time_interval
 
-from . import const
+from . import const, discovery
 from .cluster_handlers import ClusterHandler, ZDOClusterHandler
 from .const import (
     ATTR_ACTIVE_COORDINATOR,
@@ -432,6 +432,7 @@ class ZHADevice(LogMixin):
                 zha_dev.async_update_sw_build_id,
             )
         )
+        discovery.PROBE.discover_device_entities(zha_dev)
         return zha_dev
 
     @callback

homeassistant/components/zha/core/discovery.py

@@ -6,6 +6,8 @@ from collections.abc import Callable
 import logging
 from typing import TYPE_CHECKING, cast
 
+from slugify import slugify
+from zigpy.state import State
 from zigpy.zcl.clusters.general import Ota
 
 from homeassistant.const import CONF_TYPE, Platform
@@ -104,6 +106,52 @@ class ProbeEndpoint:
         self.discover_multi_entities(endpoint, config_diagnostic_entities=True)
         zha_regs.ZHA_ENTITIES.clean_up()
 
+    @callback
+    def discover_device_entities(self, device: ZHADevice) -> None:
+        """Discover entities for a ZHA device."""
+        _LOGGER.debug(
+            "Discovering entities for device: %s-%s",
+            str(device.ieee),
+            device.name,
+        )
+
+        if device.is_coordinator:
+            self.discover_coordinator_device_entities(device)
+
+    @callback
+    def discover_coordinator_device_entities(self, device: ZHADevice) -> None:
+        """Discover entities for the coordinator device."""
+        _LOGGER.debug(
+            "Discovering entities for coordinator device: %s-%s",
+            str(device.ieee),
+            device.name,
+        )
+        state: State = device.gateway.application_controller.state
+        platforms: dict[Platform, list] = get_zha_data(device.hass).platforms
+
+        @callback
+        def process_counters(counter_groups: str) -> None:
+            for counter_group, counters in getattr(state, counter_groups).items():
+                for counter in counters:
+                    platforms[Platform.SENSOR].append(
+                        (
+                            sensor.DeviceCounterSensor,
+                            (
+                                f"{slugify(str(device.ieee))}_{counter_groups}_{counter_group}_{counter}",
+                                device,
+                                counter_groups,
+                                counter_group,
+                                counter,
+                            ),
+                        )
+                    )
+
+        process_counters("counters")
+        process_counters("broadcast_counters")
+        process_counters("device_counters")
+        process_counters("group_counters")
+        zha_regs.ZHA_ENTITIES.clean_up()
+
     @callback
     def discover_by_device_type(self, endpoint: Endpoint) -> None:
         """Process an endpoint on a zigpy device."""

homeassistant/components/zha/sensor.py

@@ -1,6 +1,7 @@
 """Sensors on Zigbee Home Automation networks."""
 from __future__ import annotations
 
+import asyncio
 from dataclasses import dataclass
 from datetime import timedelta
 import enum
@@ -10,6 +11,7 @@ import random
 from typing import TYPE_CHECKING, Any, Self
 
 from zigpy import types
+from zigpy.state import Counter, State
 from zigpy.zcl.clusters.closures import WindowCovering
 from zigpy.zcl.clusters.general import Basic
 
@@ -71,7 +73,7 @@ from .core.const import (
 )
 from .core.helpers import get_zha_data
 from .core.registries import SMARTTHINGS_HUMIDITY_CLUSTER, ZHA_ENTITIES
-from .entity import ZhaEntity
+from .entity import BaseZhaEntity, ZhaEntity
 
 if TYPE_CHECKING:
     from .core.cluster_handlers import ClusterHandler
@@ -244,6 +246,83 @@ class PollableSensor(Sensor):
             )
 
 
+class DeviceCounterSensor(BaseZhaEntity, SensorEntity):
+    """Device counter sensor."""
+
+    _attr_should_poll = True
+    _attr_state_class: SensorStateClass = SensorStateClass.TOTAL
+    _attr_entity_category = EntityCategory.DIAGNOSTIC
+    _attr_entity_registry_enabled_default = False
+
+    @classmethod
+    def create_entity(
+        cls,
+        unique_id: str,
+        zha_device: ZHADevice,
+        counter_groups: str,
+        counter_group: str,
+        counter: str,
+        **kwargs: Any,
+    ) -> Self | None:
+        """Entity Factory.
+
+        Return entity if it is a supported configuration, otherwise return None
+        """
+        return cls(
+            unique_id, zha_device, counter_groups, counter_group, counter, **kwargs
+        )
+
+    def __init__(
+        self,
+        unique_id: str,
+        zha_device: ZHADevice,
+        counter_groups: str,
+        counter_group: str,
+        counter: str,
+        **kwargs: Any,
+    ) -> None:
+        """Init this sensor."""
+        super().__init__(unique_id, zha_device, **kwargs)
+        state: State = self._zha_device.gateway.application_controller.state
+        self._zigpy_counter: Counter = (
+            getattr(state, counter_groups).get(counter_group, {}).get(counter, None)
+        )
+        self._attr_name: str = self._zigpy_counter.name
+        self.remove_future: asyncio.Future
+
+    @property
+    def available(self) -> bool:
+        """Return entity availability."""
+        return self._zha_device.available
+
+    async def async_added_to_hass(self) -> None:
+        """Run when about to be added to hass."""
+        self.remove_future = self.hass.loop.create_future()
+        self._zha_device.gateway.register_entity_reference(
+            self._zha_device.ieee,
+            self.entity_id,
+            self._zha_device,
+            {},
+            self.device_info,
+            self.remove_future,
+        )
+
+    async def async_will_remove_from_hass(self) -> None:
+        """Disconnect entity object when removed."""
+        await super().async_will_remove_from_hass()
+        self.zha_device.gateway.remove_entity_reference(self)
+        self.remove_future.set_result(True)
+
+    @property
+    def native_value(self) -> StateType:
+        """Return the state of the entity."""
+        return self._zigpy_counter.value
+
+    async def async_update(self) -> None:
+        """Retrieve latest state."""
+        self.async_write_ha_state()
+
+
 # pylint: disable-next=hass-invalid-inheritance # needs fixing
 class EnumSensor(Sensor):
     """Sensor with value from enum."""

tests/components/zha/conftest.py

@@ -38,6 +38,7 @@ from tests.components.light.conftest import mock_light_profiles  # noqa: F401
 
 FIXTURE_GRP_ID = 0x1001
 FIXTURE_GRP_NAME = "fixture group"
+COUNTER_NAMES = ["counter_1", "counter_2", "counter_3"]
 
 
 @pytest.fixture(scope="session", autouse=True)
@@ -173,6 +174,10 @@ async def zigpy_app_controller():
     app.state.network_info.extended_pan_id = app.state.node_info.ieee
     app.state.network_info.channel = 15
     app.state.network_info.network_key.key = zigpy.types.KeyData(range(16))
+    app.state.counters = zigpy.state.CounterGroups()
+    app.state.counters["ezsp_counters"] = zigpy.state.CounterGroup("ezsp_counters")
+    for name in COUNTER_NAMES:
+        app.state.counters["ezsp_counters"][name].increment()
 
     # Create a fake coordinator device
     dev = app.add_device(nwk=app.state.node_info.nwk, ieee=app.state.node_info.ieee)

tests/components/zha/test_discover.py

@@ -48,6 +48,7 @@ IGNORE_SUFFIXES = [
     "off_transition_time",
     "default_move_rate",
     "start_up_current_level",
+    "counter",
 ]
 
 

tests/components/zha/test_sensor.py

@@ -9,6 +9,7 @@ from zigpy.zcl.clusters import general, homeautomation, hvac, measurement, smart
 from zigpy.zcl.clusters.hvac import Thermostat
 
 from homeassistant.components.sensor import SensorDeviceClass
+from homeassistant.components.zha.core import ZHADevice
 from homeassistant.components.zha.core.const import ZHA_CLUSTER_HANDLER_READS_PER_REQ
 import homeassistant.config as config_util
 from homeassistant.const import (
@@ -32,7 +33,7 @@ from homeassistant.const import (
     UnitOfVolume,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import restore_state
+from homeassistant.helpers import entity_registry as er, restore_state
 from homeassistant.helpers.entity_component import async_update_entity
 from homeassistant.util import dt as dt_util
 
@@ -47,6 +48,7 @@ from .common import (
 from .conftest import SIG_EP_INPUT, SIG_EP_OUTPUT, SIG_EP_PROFILE, SIG_EP_TYPE
 
 from tests.common import (
+    MockConfigEntry,
     async_fire_time_changed,
     async_mock_load_restore_state_from_storage,
 )
@@ -1183,3 +1185,60 @@ async def test_elec_measurement_skip_unsupported_attribute(
         a for call in cluster.read_attributes.call_args_list for a in call[0][0]
     }
     assert read_attrs == supported_attributes
+
+
+@pytest.fixture
+async def coordinator(hass: HomeAssistant, zigpy_device_mock, zha_device_joined):
+    """Test ZHA fan platform."""
+
+    zigpy_device = zigpy_device_mock(
+        {
+            1: {
+                SIG_EP_INPUT: [general.Groups.cluster_id],
+                SIG_EP_OUTPUT: [],
+                SIG_EP_TYPE: zigpy.profiles.zha.DeviceType.CONTROL_BRIDGE,
+                SIG_EP_PROFILE: zigpy.profiles.zha.PROFILE_ID,
+            }
+        },
+        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.available = True
+    return zha_device
+
+
+async def test_device_counter_sensors(
+    hass: HomeAssistant,
+    coordinator: ZHADevice,
+    entity_registry: er.EntityRegistry,
+    config_entry: MockConfigEntry,
+) -> None:
+    """Test quirks defined sensor."""
+
+    entity_id = "sensor.coordinator_manufacturer_coordinator_model_counter_1"
+    state = hass.states.get(entity_id)
+    assert state is None
+
+    # Enable the entity.
+    entity_registry.async_update_entity(entity_id, disabled_by=None)
+    await hass.config_entries.async_reload(config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.state == "1"
+
+    # simulate counter increment on application
+    coordinator.device.application.state.counters["ezsp_counters"][
+        "counter_1"
+    ].increment()
+
+    next_update = dt_util.utcnow() + timedelta(seconds=60)
+    async_fire_time_changed(hass, next_update)
+    await hass.async_block_till_done()
+
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.state == "2"