diff --git a/tests/test_enocean.py b/tests/test_enocean.py new file mode 100644 index 0000000000..1ff2deef33 --- /dev/null +++ b/tests/test_enocean.py @@ -0,0 +1,163 @@ +"""Tests for the EnOcean quirks.""" + +from unittest.mock import MagicMock + +import pytest +from zha.application.platforms.event import EntityEventTriggeredEvent +from zha.application.platforms.event.const import ButtonEventType +from zha.quirks import DEVICE_REGISTRY +from zigpy.device import ( + GreenPowerCommandReceived, + GreenPowerDevice as ZigpyGreenPowerDevice, +) +from zigpy.zgp.commands import GPContactStatusPayload +from zigpy.zgp.types import ApplicationID, GPDCommandID, SrcID + +import zhaquirks +from zhaquirks.enocean.ptm216z import PTM216ZButton + +zhaquirks.setup() + +# Every press captured from a real PTM 216Z: the four contacts individually, both +# contacts of each rocker, and both diagonals +CAPTURED_PRESSES = [0x01, 0x04, 0x02, 0x08, 0x03, 0x0C, 0x05, 0x0A] + + +@pytest.fixture +def ptm216z(MockAppController): + """Create the zigpy GPD for a PTM 216Z.""" + return ZigpyGreenPowerDevice( + MockAppController, + application_id=ApplicationID.SrcID, + src_id=SrcID(0x0155F443), + ) + + +@pytest.fixture +def buttons(ptm216z): + """Create the quirk's event entities, keyed by their unique ID suffix.""" + entry = DEVICE_REGISTRY.match_green_power_entry(ptm216z) + zha_device = entry.zha_device_factory(ptm216z, MagicMock()) + + entities = [ + entity + for entity in zha_device.discover_entities() + if isinstance(entity, PTM216ZButton) + ] + + for entity in entities: + entity.on_add() + + return {entity.unique_id.rsplit("-", 1)[-1]: entity for entity in entities} + + +def emit(device, command_id: GPDCommandID, contact_status: int) -> None: + """Emit a GPD 8-bit vector command.""" + device.emit( + GreenPowerCommandReceived.event_type, + GreenPowerCommandReceived( + device_ieee=str(device.ieee), + endpoint_id=None, + command_id=command_id, + command=GPContactStatusPayload(contact_status=contact_status), + ), + ) + + +def test_ptm216z_matches(ptm216z): + """Test that the quirk matches a PTM 216Z by its SrcID.""" + entry = DEVICE_REGISTRY.match_green_power_entry(ptm216z) + + assert entry is not None + assert entry.source.module == "zhaquirks.enocean.ptm216z" + + +def test_ptm216z_does_not_match_other_vendors(MockAppController): + """Test that the quirk ignores a GPD outside of EnOcean's SrcID block.""" + gpd = ZigpyGreenPowerDevice( + MockAppController, + application_id=ApplicationID.SrcID, + src_id=SrcID(0x01700001), + ) + + assert DEVICE_REGISTRY.match_green_power_entry(gpd) is None + + +def test_ptm216z_entities(buttons): + """Test that there is one entity per contact.""" + assert list(buttons) == ["button_a0", "button_a1", "button_b0", "button_b1"] + assert [entity.event_types for entity in buttons.values()] == [ + [ButtonEventType.PRESS_START, ButtonEventType.PRESS_END] + ] * 4 + assert buttons["button_a0"].primary + assert not any(entity.primary for entity in list(buttons.values())[1:]) + + +@pytest.mark.parametrize( + ("contact_status", "expected"), + [ + (0x01, ["button_a0"]), + (0x02, ["button_a1"]), + (0x04, ["button_b0"]), + (0x08, ["button_b1"]), + # Both contacts of a rocker, and the two diagonals + (0x03, ["button_a0", "button_a1"]), + (0x0C, ["button_b0", "button_b1"]), + (0x05, ["button_a0", "button_b0"]), + (0x0A, ["button_a1", "button_b1"]), + ], +) +def test_ptm216z_press_and_release(ptm216z, buttons, contact_status, expected): + """Test that a press and its release only fire the contacts that were pressed.""" + events: list[EntityEventTriggeredEvent] = [] + for entity in buttons.values(): + entity.on_event(EntityEventTriggeredEvent.event, events.append) + + emit(ptm216z, GPDCommandID.Press8BitVector, contact_status) + assert [event.unique_id.rsplit("-", 1)[-1] for event in events] == expected + assert all( + event.triggered.event_type == ButtonEventType.PRESS_START for event in events + ) + assert events[0].triggered.event_attributes == {"contact_status": contact_status} + + events.clear() + + # The switch only reports that every contact is open again + emit(ptm216z, GPDCommandID.Release8BitVector, 0) + assert [event.unique_id.rsplit("-", 1)[-1] for event in events] == expected + assert all( + event.triggered.event_type == ButtonEventType.PRESS_END for event in events + ) + + +def test_ptm216z_release_without_press(ptm216z, buttons): + """Test that a release nothing was pressed for fires nothing.""" + events: list[EntityEventTriggeredEvent] = [] + for entity in buttons.values(): + entity.on_event(EntityEventTriggeredEvent.event, events.append) + + emit(ptm216z, GPDCommandID.Release8BitVector, 0) + + assert events == [] + + +def test_ptm216z_captured_sequence(ptm216z, buttons): + """Test the full captured sequence of presses, each followed by its release.""" + events: list[EntityEventTriggeredEvent] = [] + for entity in buttons.values(): + entity.on_event(EntityEventTriggeredEvent.event, events.append) + + for contact_status in CAPTURED_PRESSES: + emit(ptm216z, GPDCommandID.Press8BitVector, contact_status) + emit(ptm216z, GPDCommandID.Release8BitVector, 0) + + presses = [ + e for e in events if e.triggered.event_type == ButtonEventType.PRESS_START + ] + releases = [ + e for e in events if e.triggered.event_type == ButtonEventType.PRESS_END + ] + + # Four single-contact presses and four two-contact presses, each released once + assert len(presses) == 4 + 2 * 4 + assert [e.unique_id for e in presses] == [e.unique_id for e in releases] diff --git a/tests/test_green_power.py b/tests/test_green_power.py new file mode 100644 index 0000000000..60553e1a60 --- /dev/null +++ b/tests/test_green_power.py @@ -0,0 +1,275 @@ +"""Tests for the Green Power quirk builder.""" + +from typing import Any +from unittest.mock import MagicMock + +import pytest +from zha.application.platforms.event import ( + BaseEvent, + EntityEventTriggeredEvent, + TriggeredEvent, +) +from zha.quirks import DEVICE_REGISTRY, QUIRK_REGISTRY_ENTRY_ATTR +from zha.zigbee.device import GreenPowerDevice +from zigpy.device import ( + GreenPowerCommandReceived, + GreenPowerDevice as ZigpyGreenPowerDevice, +) +from zigpy.types import EUI64 +from zigpy.zgp.types import ApplicationID, DeviceID, GPDCommandID, SrcID + +from zhaquirks.builder import EventDeviceClass, GreenPowerQuirkBuilder +from zhaquirks.builder.green_power import QuirkGreenPowerDevice + + +class CommandEvent(BaseEvent): + """Event entity firing an event type for a single GPD command.""" + + _attr_device_class = EventDeviceClass.BUTTON + + def __init__( + self, + device: GreenPowerDevice, + *, + command_id: GPDCommandID, + event_type: str, + **kwargs: Any, + ) -> None: + """Initialize the event entity.""" + self._command_id = command_id + self._attr_event_types = [event_type] + + super().__init__(device, **kwargs) + + def on_add(self) -> None: + """Subscribe to the commands sent by the device.""" + super().on_add() + self._on_remove_callbacks.append( + self.device.device.on_event( + GreenPowerCommandReceived.event_type, self._handle_gp_command + ) + ) + + def _handle_gp_command(self, event: GreenPowerCommandReceived) -> None: + """Fire the entity's event type for its own command.""" + if event.command_id == self._command_id: + self._trigger_event(self._attr_event_types[0]) + + +@pytest.fixture +def zigpy_gpd(MockAppController): + """Create a commissioned SrcID-addressed zigpy GPD.""" + gpd = ZigpyGreenPowerDevice( + MockAppController, + application_id=ApplicationID.SrcID, + src_id=SrcID(0x01700001), + ) + gpd.device_id = DeviceID.OnOffSwitch + gpd.commands = [ + GPDCommandID.Toggle, + GPDCommandID.Press1of1, + GPDCommandID.Release1of1, + ] + return gpd + + +def test_green_power_quirk_builder(zigpy_gpd): + """Test matching and resolving a Green Power quirk.""" + with DEVICE_REGISTRY.preserve_state(): + entry = ( + GreenPowerQuirkBuilder() + .applies_to(device_id=DeviceID.OnOffSwitch) + .src_id_range(0x01700000, 0x0170FFFF) + .filter(lambda device: GPDCommandID.Toggle in device.commands) + .friendly_name(manufacturer="EnOcean", model="PTM 215Z") + .add_to_registry() + ) + + assert DEVICE_REGISTRY.match_green_power_entry(zigpy_gpd) is entry + + resolved = DEVICE_REGISTRY.resolve(zigpy_gpd) + assert resolved is zigpy_gpd + assert getattr(resolved, QUIRK_REGISTRY_ENTRY_ATTR) is entry + + zha_device = entry.zha_device_factory(resolved, MagicMock()) + assert isinstance(zha_device, QuirkGreenPowerDevice) + assert zha_device.manufacturer == "EnOcean" + assert zha_device.model == "PTM 215Z" + assert zha_device.name == "EnOcean PTM 215Z" + + +def test_green_power_quirk_builder_no_match(zigpy_gpd): + """Test that a non-matching quirk leaves the device unquirked.""" + with DEVICE_REGISTRY.preserve_state(): + ( + GreenPowerQuirkBuilder() + .applies_to(device_id=DeviceID.GenericSwitch) + .add_to_registry() + ) + ( + GreenPowerQuirkBuilder() + .applies_to(device_id=DeviceID.OnOffSwitch) + .src_id_range(0x00400000, 0x0040FFFF) + .add_to_registry() + ) + + assert DEVICE_REGISTRY.match_green_power_entry(zigpy_gpd) is None + assert DEVICE_REGISTRY.resolve(zigpy_gpd) is zigpy_gpd + assert not hasattr(zigpy_gpd, QUIRK_REGISTRY_ENTRY_ATTR) + + +def test_green_power_quirk_builder_ieee_prefix(MockAppController): + """Test matching an IEEE-addressed GPD by address prefix.""" + gpd = ZigpyGreenPowerDevice( + MockAppController, + application_id=ApplicationID.IEEE, + ieee=EUI64.convert("04:cd:15:00:11:22:33:44"), + ) + + with DEVICE_REGISTRY.preserve_state(): + entry = GreenPowerQuirkBuilder().ieee_prefix("04:cd:15").add_to_registry() + + assert DEVICE_REGISTRY.match_green_power_entry(gpd) is entry + + +def test_green_power_quirk_builder_requires_criteria(): + """Test that a quirk without matching criteria is rejected.""" + with pytest.raises(ValueError): + GreenPowerQuirkBuilder().friendly_name( + manufacturer="Acme", model="Switch" + ).add_to_registry() + + with pytest.raises(ValueError): + GreenPowerQuirkBuilder().applies_to() + + +def test_green_power_quirk_custom_device_class(zigpy_gpd): + """Test that a custom ZHA device class is used by the factory.""" + + class CustomGreenPowerDevice(QuirkGreenPowerDevice): + """Custom quirk device class.""" + + with DEVICE_REGISTRY.preserve_state(): + entry = ( + GreenPowerQuirkBuilder() + .applies_to(device_id=DeviceID.OnOffSwitch) + .device_class(CustomGreenPowerDevice) + .add_to_registry() + ) + + zha_device = entry.zha_device_factory(zigpy_gpd, MagicMock()) + assert isinstance(zha_device, CustomGreenPowerDevice) + + +def test_green_power_quirk_entities(zigpy_gpd): + """Test entities declared by a Green Power quirk.""" + with DEVICE_REGISTRY.preserve_state(): + entry = ( + GreenPowerQuirkBuilder() + .applies_to(device_id=DeviceID.OnOffSwitch) + .entity( + CommandEvent, + command_id=GPDCommandID.Press1of1, + event_type="press", + unique_id_suffix="button_1", + primary=True, + fallback_name="Button 1", + ) + .entity( + CommandEvent, + command_id=GPDCommandID.Toggle, + event_type="toggle", + unique_id_suffix="button_2", + fallback_name="Button 2", + ) + .add_to_registry() + ) + + zha_device = entry.zha_device_factory(zigpy_gpd, MagicMock()) + entities = list(zha_device.discover_entities()) + assert len(entities) == 2 + + button_1, button_2 = entities + assert button_1.unique_id == f"{zigpy_gpd.ieee}-button_1" + assert button_1.event_types == ["press"] + assert button_1.device_class == EventDeviceClass.BUTTON + assert button_1.fallback_name == "Button 1" + assert button_1.primary + assert button_2.unique_id == f"{zigpy_gpd.ieee}-button_2" + assert button_2.event_types == ["toggle"] + assert not button_2.primary + + events: list[EntityEventTriggeredEvent] = [] + for entity in entities: + entity.on_add() + entity.on_event(EntityEventTriggeredEvent.event, events.append) + + # A press fires only the first entity + zigpy_gpd.emit( + GreenPowerCommandReceived.event_type, + GreenPowerCommandReceived( + device_ieee=str(zigpy_gpd.ieee), + endpoint_id=0, + command_id=GPDCommandID.Press1of1, + command=None, + ), + ) + assert [(event.unique_id, event.triggered) for event in events] == [ + ( + button_1.unique_id, + TriggeredEvent(event_type="press", event_attributes={}), + ) + ] + + # A toggle fires only the second + zigpy_gpd.emit( + GreenPowerCommandReceived.event_type, + GreenPowerCommandReceived( + device_ieee=str(zigpy_gpd.ieee), + endpoint_id=0, + command_id=GPDCommandID.Toggle, + command=None, + ), + ) + assert events[-1].unique_id == button_2.unique_id + assert events[-1].triggered == TriggeredEvent( + event_type="toggle", event_attributes={} + ) + + +async def test_green_power_quirk_entity_removed(zigpy_gpd): + """Test that a removed entity stops listening for GPD commands.""" + with DEVICE_REGISTRY.preserve_state(): + entry = ( + GreenPowerQuirkBuilder() + .applies_to(device_id=DeviceID.OnOffSwitch) + .entity( + CommandEvent, + command_id=GPDCommandID.Toggle, + event_type="toggle", + fallback_name="Button", + ) + .add_to_registry() + ) + + zha_device = entry.zha_device_factory(zigpy_gpd, MagicMock()) + (entity,) = zha_device.discover_entities() + entity.on_add() + + events: list[EntityEventTriggeredEvent] = [] + entity.on_event(EntityEventTriggeredEvent.event, events.append) + + toggle = GreenPowerCommandReceived( + device_ieee=str(zigpy_gpd.ieee), + endpoint_id=0, + command_id=GPDCommandID.Toggle, + command=None, + ) + + zigpy_gpd.emit(GreenPowerCommandReceived.event_type, toggle) + assert len(events) == 1 + + await entity.on_remove() + + zigpy_gpd.emit(GreenPowerCommandReceived.event_type, toggle) + assert len(events) == 1 diff --git a/zhaquirks/builder/__init__.py b/zhaquirks/builder/__init__.py index 04e557c46b..81e5036e43 100644 --- a/zhaquirks/builder/__init__.py +++ b/zhaquirks/builder/__init__.py @@ -9,6 +9,7 @@ from zha.application import EntityPlatform, EntityType from zha.application.platforms.binary_sensor.device_class import BinarySensorDeviceClass +from zha.application.platforms.event.const import ButtonEventType, EventDeviceClass from zha.application.platforms.number.device_class import NumberDeviceClass from zha.application.platforms.sensor.device_class import ( SensorDeviceClass, @@ -17,12 +18,16 @@ from zha.units import * # noqa: F401, F403 from zhaquirks.builder.builder import QuirkBuilder +from zhaquirks.builder.green_power import GreenPowerQuirkBuilder from zhaquirks.builder.metadata import ReportingConfig __all__ = [ "BinarySensorDeviceClass", + "ButtonEventType", "EntityPlatform", "EntityType", + "EventDeviceClass", + "GreenPowerQuirkBuilder", "NumberDeviceClass", "QuirkBuilder", "ReportingConfig", diff --git a/zhaquirks/builder/green_power.py b/zhaquirks/builder/green_power.py new file mode 100644 index 0000000000..268ee91396 --- /dev/null +++ b/zhaquirks/builder/green_power.py @@ -0,0 +1,225 @@ +"""The declarative (GreenPowerQuirkBuilder) authoring API for Green Power devices.""" + +from __future__ import annotations + +from collections.abc import Iterator +from dataclasses import dataclass +import inspect +import pathlib +from types import FrameType +from typing import TYPE_CHECKING, Any, Self + +from frozendict import frozendict + +# `discovery` must be imported before `zha.zigbee.device`: the platform modules +# it loads participate in an import cycle with the device module and cannot be +# loaded while `zha.zigbee.device` is only partially initialized. +from zha.application import discovery # noqa: F401 +from zha.application.platforms import BaseEntity, PlatformEntity +from zha.quirks import ( + DEVICE_REGISTRY, + DeviceRegistry, + GreenPowerDeviceMatch, + GreenPowerFilterType, + GreenPowerQuirkRegistryEntry, + QuirkSource, +) +from zha.zigbee.device import GreenPowerDevice +import zigpy.device +from zigpy.zgp.types import DeviceID, SrcID + +if TYPE_CHECKING: + from zha.application.gateway import Gateway + + +@dataclass(frozen=True) +class GreenPowerQuirkDefinition: + """ZHA-level metadata for a Green Power quirk.""" + + friendly_manufacturer: str | None = None + friendly_model: str | None = None + entities: tuple[tuple[type[PlatformEntity], frozendict[str, Any]], ...] = () + + +class QuirkGreenPowerDevice(GreenPowerDevice): + """Base ZHA device for GreenPowerQuirkBuilder.""" + + def __init__( + self, + zigpy_device: zigpy.device.GreenPowerDevice, + gateway: Gateway, + *, + quirk_definition: GreenPowerQuirkDefinition, + ) -> None: + """Initialize the quirk device.""" + self._quirk_definition = quirk_definition + super().__init__(zigpy_device, gateway) + + @property + def quirk_metadata(self) -> GreenPowerQuirkDefinition: + """Return the ZHA-level quirk metadata for this device.""" + return self._quirk_definition + + def discover_entities(self) -> Iterator[BaseEntity]: + """Yield the quirk's entities.""" + yield from super().discover_entities() + + for entity_class, kwargs in self._quirk_definition.entities: + yield entity_class( + self, unique_id=str(self.ieee), from_quirk=True, **kwargs + ) + + def _resolve_manufacturer(self) -> str: + if self._quirk_definition.friendly_manufacturer is not None: + return self._quirk_definition.friendly_manufacturer + return super()._resolve_manufacturer() + + def _resolve_model(self) -> str: + if self._quirk_definition.friendly_model is not None: + return self._quirk_definition.friendly_model + return super()._resolve_model() + + +@dataclass(frozen=True) +class GreenPowerQuirkFactory: + """Registry-entry factory building a `QuirkGreenPowerDevice` bound to its definition.""" + + base: type[QuirkGreenPowerDevice] + quirk_definition: GreenPowerQuirkDefinition + + def __call__( + self, zigpy_device: zigpy.device.GreenPowerDevice, gateway: Gateway + ) -> QuirkGreenPowerDevice: + """Build the bound `QuirkGreenPowerDevice` for a resolved zigpy device.""" + return self.base(zigpy_device, gateway, quirk_definition=self.quirk_definition) + + +class GreenPowerQuirkBuilder: + """Builder compiling a declarative Green Power quirk into a registry entry.""" + + def __init__(self, registry: DeviceRegistry = DEVICE_REGISTRY) -> None: + """Initialize the quirk builder.""" + self.registry: DeviceRegistry = registry + self.device_id: DeviceID | None = None + self.manufacturer_id: int | None = None + self.model_id: int | None = None + self.src_id_ranges: list[tuple[SrcID, SrcID]] = [] + self.ieee_prefixes: list[bytes] = [] + self.filters: list[GreenPowerFilterType] = [] + self.friendly_manufacturer: str | None = None + self.friendly_model: str | None = None + self.entities: list[tuple[type[PlatformEntity], frozendict[str, Any]]] = [] + self.custom_device_class: type[QuirkGreenPowerDevice] | None = None + + current_frame: FrameType = inspect.currentframe() + caller: FrameType = current_frame.f_back + self.quirk_file = pathlib.Path(caller.f_code.co_filename) + self.quirk_file_line = caller.f_lineno + self.quirk_module: str = caller.f_globals["__name__"] + + def applies_to( + self, + *, + device_id: DeviceID | int | None = None, + manufacturer_id: int | None = None, + model_id: int | None = None, + ) -> Self: + """Match the GPD's commissioning identity fields.""" + if device_id is None and manufacturer_id is None and model_id is None: + raise ValueError( + "At least one of device_id, manufacturer_id, or model_id must be" + " specified" + ) + + if device_id is not None: + self.device_id = DeviceID(device_id) + if manufacturer_id is not None: + self.manufacturer_id = manufacturer_id + if model_id is not None: + self.model_id = model_id + + return self + + def src_id_range(self, lower: int, upper: int) -> Self: + """Match SrcIDs within `[lower, upper]`: vendors allocate SrcID blocks.""" + self.src_id_ranges.append((SrcID(lower), SrcID(upper))) + return self + + def ieee_prefix(self, prefix: str) -> Self: + """Match IEEE-addressed GPDs whose address begins with `prefix` ("04:cd:15").""" + self.ieee_prefixes.append(bytes(int(octet, 16) for octet in prefix.split(":"))) + return self + + def filter(self, filter_function: GreenPowerFilterType) -> Self: + """Add a filter and return self. + + The filter function should take a single argument, a + `zigpy.device.GreenPowerDevice` instance, and return a boolean if the + condition the filter is testing passes. + """ + self.filters.append(filter_function) + return self + + def friendly_name(self, *, manufacturer: str, model: str) -> Self: + """Override the device name displayed in HA.""" + self.friendly_manufacturer = manufacturer + self.friendly_model = model + return self + + def entity(self, entity_class: type[PlatformEntity], **kwargs: Any) -> Self: + """Add an entity, passing `kwargs` to its constructor.""" + self.entities.append((entity_class, frozendict(kwargs))) + return self + + def device_class(self, custom_device_class: type[QuirkGreenPowerDevice]) -> Self: + """Set a custom ZHA device class.""" + assert issubclass(custom_device_class, QuirkGreenPowerDevice) + self.custom_device_class = custom_device_class + return self + + def add_to_registry( + self, registry: DeviceRegistry | None = None + ) -> GreenPowerQuirkRegistryEntry: + """Compile the quirk into a `GreenPowerQuirkRegistryEntry` and register it.""" + device_match = GreenPowerDeviceMatch( + device_id=self.device_id, + manufacturer_id=self.manufacturer_id, + model_id=self.model_id, + src_id_ranges=tuple(self.src_id_ranges), + ieee_prefixes=tuple(self.ieee_prefixes), + filters=tuple(self.filters), + ) + + if device_match == GreenPowerDeviceMatch(): + raise ValueError( + "At least one matching criterion must be specified for a Green Power" + " quirk" + ) + + quirk_definition = GreenPowerQuirkDefinition( + friendly_manufacturer=self.friendly_manufacturer, + friendly_model=self.friendly_model, + entities=tuple(self.entities), + ) + + base = ( + self.custom_device_class + if self.custom_device_class + else QuirkGreenPowerDevice + ) + zha_device_factory = GreenPowerQuirkFactory(base, quirk_definition) + + entry = GreenPowerQuirkRegistryEntry( + device_match=device_match, + zha_device_factory=zha_device_factory, + source=QuirkSource( + module=self.quirk_module, + file=str(self.quirk_file), + line=self.quirk_file_line, + label=f"({self.friendly_manufacturer} / {self.friendly_model})", + ), + ) + + (registry or self.registry).register(entry) + + return entry diff --git a/zhaquirks/enocean/__init__.py b/zhaquirks/enocean/__init__.py new file mode 100644 index 0000000000..9a1c23e688 --- /dev/null +++ b/zhaquirks/enocean/__init__.py @@ -0,0 +1 @@ +"""Module for EnOcean quirks implementations.""" diff --git a/zhaquirks/enocean/ptm216z.py b/zhaquirks/enocean/ptm216z.py new file mode 100644 index 0000000000..f0db3fe4d1 --- /dev/null +++ b/zhaquirks/enocean/ptm216z.py @@ -0,0 +1,98 @@ +"""EnOcean PTM 216Z self-powered double rocker switch.""" + +from typing import Any + +from zha.application.platforms.event import BaseEvent +from zha.zigbee.device import GreenPowerDevice +from zigpy.device import GreenPowerCommandReceived +from zigpy.zgp.types import GPDCommandID + +from zhaquirks.builder import ButtonEventType, EventDeviceClass, GreenPowerQuirkBuilder + +# The switch has two rockers, each with two contacts, reported as a bit in the 8-bit +# vector of the GPD press command. Bits 4-7 are unused. +CONTACT_A0 = 0b0001 +CONTACT_A1 = 0b0010 +CONTACT_B0 = 0b0100 +CONTACT_B1 = 0b1000 + + +class PTM216ZButton(BaseEvent): + """Event entity for a single contact of the switch.""" + + _attr_device_class = EventDeviceClass.BUTTON + _attr_event_types = [ButtonEventType.PRESS_START, ButtonEventType.PRESS_END] + + def __init__( + self, device: GreenPowerDevice, *, contact: int, **kwargs: Any + ) -> None: + """Initialize the event entity for a single contact.""" + self._contact = contact + self._pressed = False + + super().__init__(device, **kwargs) + + def on_add(self) -> None: + """Subscribe to the commands sent by the switch.""" + super().on_add() + self._on_remove_callbacks.append( + self.device.device.on_event( + GreenPowerCommandReceived.event_type, self._handle_contact_status + ) + ) + + def _handle_contact_status(self, event: GreenPowerCommandReceived) -> None: + """Fire press and release events for this entity's own contact.""" + # zigpy reports a command it could not parse without its payload + if event.command is None: + return + + if event.command_id == GPDCommandID.Press8BitVector: + pressed = bool(event.command.contact_status & self._contact) + elif event.command_id == GPDCommandID.Release8BitVector: + # The switch only reports that every contact is open again + pressed = False + else: + return + + if pressed == self._pressed: + return + + self._pressed = pressed + self._trigger_event( + ButtonEventType.PRESS_START if pressed else ButtonEventType.PRESS_END, + event.command.as_dict(), + ) + + +( + GreenPowerQuirkBuilder() + .src_id_range(0x01550000, 0x0155FFFF) + .friendly_name(manufacturer="EnOcean", model="PTM 216Z") + .entity( + PTM216ZButton, + contact=CONTACT_A0, + unique_id_suffix="button_a0", + fallback_name="Rocker A top", + primary=True, + ) + .entity( + PTM216ZButton, + contact=CONTACT_A1, + unique_id_suffix="button_a1", + fallback_name="Rocker A bottom", + ) + .entity( + PTM216ZButton, + contact=CONTACT_B0, + unique_id_suffix="button_b0", + fallback_name="Rocker B top", + ) + .entity( + PTM216ZButton, + contact=CONTACT_B1, + unique_id_suffix="button_b1", + fallback_name="Rocker B bottom", + ) + .add_to_registry() +)