Add lqi and rssi sensors back to ZHA (#62716)

* update device list

* Only 1 identify button per device

* Add LQI and RSSI sensors to ZHA

* refactor entity creation filter

* update device list and update discover test

* fix reference

* code reduction

* walrus

* parens

* simplify

homeassistant/components/zha/button.py

@@ -89,11 +89,10 @@ class ZHAIdentifyButton(ZHAButton):
 
         Return entity if it is a supported configuration, otherwise return None
         """
-        platform_restrictions = ZHA_ENTITIES.single_device_matches[Platform.BUTTON]
+        if ZHA_ENTITIES.prevent_entity_creation(
-        device_restrictions = platform_restrictions[zha_device.ieee]
+            Platform.BUTTON, zha_device.ieee, CHANNEL_IDENTIFY
-        if CHANNEL_IDENTIFY in device_restrictions:
+        ):
             return None
-        device_restrictions.append(CHANNEL_IDENTIFY)
         return cls(unique_id, zha_device, channels, **kwargs)
 
     _attr_device_class: ButtonDeviceClass = ButtonDeviceClass.UPDATE

homeassistant/components/zha/core/registries.py

@@ -346,6 +346,15 @@ class ZHAEntityRegistry:
 
         return decorator
 
+    def prevent_entity_creation(self, platform: Platform, ieee: EUI64, key: str):
+        """Return True if the entity should not be created."""
+        platform_restrictions = self.single_device_matches[platform]
+        device_restrictions = platform_restrictions[ieee]
+        if key in device_restrictions:
+            return True
+        device_restrictions.append(key)
+        return False
+
     def clean_up(self) -> None:
         """Clean up post discovery."""
         self.single_device_matches: dict[

homeassistant/components/zha/entity.py

@@ -41,7 +41,7 @@ UPDATE_GROUP_FROM_CHILD_DELAY = 0.5
 class BaseZhaEntity(LogMixin, entity.Entity):
     """A base class for ZHA entities."""
 
-    _unique_id_suffix: str | None = None
+    unique_id_suffix: str | None = None
 
     def __init__(self, unique_id: str, zha_device: ZhaDeviceType, **kwargs) -> None:
         """Init ZHA entity."""
@@ -49,8 +49,8 @@ class BaseZhaEntity(LogMixin, entity.Entity):
         self._force_update: bool = False
         self._should_poll: bool = False
         self._unique_id: str = unique_id
-        if self._unique_id_suffix:
+        if self.unique_id_suffix:
-            self._unique_id += f"-{self._unique_id_suffix}"
+            self._unique_id += f"-{self.unique_id_suffix}"
         self._state: Any = None
         self._extra_state_attributes: dict[str, Any] = {}
         self._zha_device: ZhaDeviceType = zha_device
@@ -154,7 +154,7 @@ class ZhaEntity(BaseZhaEntity, RestoreEntity):
         """
         super().__init_subclass__(**kwargs)
         if id_suffix:
-            cls._unique_id_suffix = id_suffix
+            cls.unique_id_suffix = id_suffix
 
     def __init__(
         self,
@@ -169,8 +169,8 @@ class ZhaEntity(BaseZhaEntity, RestoreEntity):
         ch_names = [ch.cluster.ep_attribute for ch in channels]
         ch_names = ", ".join(sorted(ch_names))
         self._name: str = f"{zha_device.name} {ieeetail} {ch_names}"
-        if self._unique_id_suffix:
+        if self.unique_id_suffix:
-            self._name += f" {self._unique_id_suffix}"
+            self._name += f" {self.unique_id_suffix}"
         self.cluster_channels: dict[str, ChannelType] = {}
         for channel in channels:
             self.cluster_channels[channel.name] = channel

homeassistant/components/zha/sensor.py

@@ -25,6 +25,7 @@ from homeassistant.const import (
     ELECTRIC_CURRENT_AMPERE,
     ELECTRIC_POTENTIAL_VOLT,
     ENERGY_KILO_WATT_HOUR,
+    ENTITY_CATEGORY_DIAGNOSTIC,
     LIGHT_LUX,
     PERCENTAGE,
     POWER_VOLT_AMPERE,
@@ -50,6 +51,7 @@ from homeassistant.helpers.typing import StateType
 from .core import discovery
 from .core.const import (
     CHANNEL_ANALOG_INPUT,
+    CHANNEL_BASIC,
     CHANNEL_ELECTRICAL_MEASUREMENT,
     CHANNEL_HUMIDITY,
     CHANNEL_ILLUMINANCE,
@@ -675,3 +677,45 @@ class SinopeHVACAction(ThermostatHVACAction):
         ):
             return CURRENT_HVAC_IDLE
         return CURRENT_HVAC_OFF
+
+
+@MULTI_MATCH(channel_names=CHANNEL_BASIC)
+class RSSISensor(Sensor, id_suffix="rssi"):
+    """RSSI sensor for a device."""
+
+    _state_class: SensorStateClass = SensorStateClass.MEASUREMENT
+    _device_class: SensorDeviceClass = SensorDeviceClass.SIGNAL_STRENGTH
+    _attr_entity_category = ENTITY_CATEGORY_DIAGNOSTIC
+    _attr_entity_registry_enabled_default = False
+
+    @classmethod
+    def create_entity(
+        cls,
+        unique_id: str,
+        zha_device: ZhaDeviceType,
+        channels: list[ChannelType],
+        **kwargs,
+    ) -> ZhaEntity | None:
+        """Entity Factory.
+
+        Return entity if it is a supported configuration, otherwise return None
+        """
+        key = f"{CHANNEL_BASIC}_{cls.unique_id_suffix}"
+        if ZHA_ENTITIES.prevent_entity_creation(Platform.SENSOR, zha_device.ieee, key):
+            return None
+        return cls(unique_id, zha_device, channels, **kwargs)
+
+    @property
+    def native_value(self) -> StateType:
+        """Return the state of the entity."""
+        return getattr(self._zha_device.device, self.unique_id_suffix)
+
+    @property
+    def should_poll(self) -> bool:
+        """Poll the entity for current state."""
+        return True
+
+
+@MULTI_MATCH(channel_names=CHANNEL_BASIC)
+class LQISensor(RSSISensor, id_suffix="lqi"):
+    """LQI sensor for a device."""

tests/components/zha/test_discover.py

@@ -2,7 +2,7 @@
 
 import re
 from unittest import mock
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
 from zigpy.const import SIG_ENDPOINTS, SIG_MANUFACTURER, SIG_MODEL, SIG_NODE_DESC
@@ -70,6 +70,14 @@ def channels_mock(zha_device_mock):
     "zigpy.zcl.clusters.general.Identify.request",
     new=AsyncMock(return_value=[mock.sentinel.data, zcl_f.Status.SUCCESS]),
 )
+# We do this here because we are testing ZHA discovery logic. Point being we want to ensure that
+# all discovered entities are dispatched for creation. In order to test this we need the entities
+# added to HA. So we ensure that they are all enabled even though they won't necessarily be in reality
+# at runtime
+@patch(
+    "homeassistant.components.zha.entity.ZhaEntity.entity_registry_enabled_default",
+    new=Mock(return_value=True),
+)
 @pytest.mark.parametrize("device", DEVICES)
 async def test_devices(
     device,