diff --git a/tests/test_device.py b/tests/test_device.py index b653ea89b..93d0cae93 100644 --- a/tests/test_device.py +++ b/tests/test_device.py @@ -16,13 +16,13 @@ ExposesFeatureMetadata, QuirkBuilder, ) -from zigpy.quirks.v2.homeassistant import EntityType +from zigpy.quirks.v2.homeassistant import PERCENTAGE, EntityType from zigpy.quirks.v2.homeassistant.sensor import SensorDeviceClass, SensorStateClass import zigpy.types from zigpy.typing import UNDEFINED from zigpy.zcl import ClusterType from zigpy.zcl.clusters import general -from zigpy.zcl.clusters.general import Ota, PowerConfiguration +from zigpy.zcl.clusters.general import AnalogInput, Ota, PowerConfiguration from zigpy.zcl.clusters.lighting import Color from zigpy.zcl.clusters.measurement import CarbonDioxideConcentration from zigpy.zcl.foundation import Status, WriteAttributesResponse @@ -1031,6 +1031,66 @@ async def test_quirks_v2_prevent_default_entities(zha_gateway: Gateway) -> None: assert len(zha_device.platform_entities) == 7 +async def test_quirks_v2_prevent_default_entities_does_not_remove_v2_entities( + zha_gateway: Gateway, +) -> None: + """Test prevent_default_entity_creation does not remove quirks v2 entities. + + `prevent_default_entity_creation` is meant to suppress default ZHA entities + only; quirks-v2-defined entities (via `.sensor(...)` etc.) must survive + even when the predicate would otherwise match them — e.g. when the v2 + entity deliberately reuses the same `unique_id_suffix` as the default + entity in order to preserve the existing HA entity registry entry. + """ + registry = DeviceRegistry() + + ( + QuirkBuilder("Espressif", "ZigbeeAnalogDevice", registry=registry) + # Both the prevent rule and the v2 sensor target the same suffix — + # without the fix this rule would remove the v2 sensor as well. + .prevent_default_entity_creation( + endpoint_id=1, + cluster_id=AnalogInput.cluster_id, + unique_id_suffix="analog_input", + ) + .sensor( + attribute_name=AnalogInput.AttributeDefs.present_value.name, + cluster_id=AnalogInput.cluster_id, + # Match the default ZHA `AnalogInputSensor._unique_id_suffix` so + # the new entity inherits the existing HA entity registry entry. + unique_id_suffix="analog_input", + device_class=SensorDeviceClass.AQI, + state_class=SensorStateClass.MEASUREMENT, + unit=PERCENTAGE, + translation_key="dirty_level", + fallback_name="Dirty Level", + ) + .add_to_registry() + ) + + zigpy_dev = registry.get_device( + await zigpy_device_from_json( + zha_gateway.application_controller, + "tests/data/devices/espressif-zigbeeanalogdevice.json", + ) + ) + + zha_device = await join_zigpy_device(zha_gateway, zigpy_dev) + + # The v2 sensor survived even though the prevent predicate (matched by + # `unique_id_suffix="analog_input"`) also matches it. + sensors = [ + e + for e in zha_device.platform_entities.values() + if e.PLATFORM == Platform.SENSOR and "analog_input" in e.unique_id + ] + assert len(sensors) == 1 + (sensor,) = sensors + assert sensor._from_quirks_v2 is True + assert sensor.translation_key == "dirty_level" + assert sensor._attr_device_class == SensorDeviceClass.AQI + + async def test_quirks_v2_change_entity_metadata(zha_gateway: Gateway) -> None: """Test quirks v2 can change entity metadata.""" registry = DeviceRegistry() diff --git a/zha/application/platforms/__init__.py b/zha/application/platforms/__init__.py index 9da3fc16c..18e730172 100644 --- a/zha/application/platforms/__init__.py +++ b/zha/application/platforms/__init__.py @@ -449,6 +449,11 @@ class PlatformEntity(BaseEntity): # entities using the same cluster handler/cluster id for the entity. _unique_id_suffix: str | None = None + # True when this entity was constructed from quirks v2 EntityMetadata. + # Used to scope filters like `prevent_default_entity_creation` to default + # ZHA entities only. + _from_quirks_v2: bool = False + _migrate_platform_unique_ids: tuple[tuple[UniqueIdMigration, str]] | None = None # Auto-discovery for the entity @@ -491,6 +496,8 @@ def __init__( def _init_from_quirks_metadata(self, entity_metadata: EntityMetadata) -> None: """Init this entity from the quirks metadata.""" + self._from_quirks_v2 = True + if entity_metadata.initially_disabled: self._attr_entity_registry_enabled_default = False diff --git a/zha/zigbee/device.py b/zha/zigbee/device.py index 14884e1da..7b8fc3cf6 100644 --- a/zha/zigbee/device.py +++ b/zha/zigbee/device.py @@ -963,6 +963,13 @@ def _is_entity_removed_by_quirk(self, entity: PlatformEntity) -> bool: if self.quirk_metadata is None: return False + # `prevent_default_entity_creation` only targets default ZHA entities. + # Entities defined by the quirk itself (via quirks v2 EntityMetadata, + # e.g. `.sensor(...)`) are not removed here, even when the predicate + # would otherwise match them. + if entity._from_quirks_v2: + return False + for meta in self.quirk_metadata.disabled_default_entities: _LOGGER.debug("Checking if entity %s is removed by %s", entity, meta)