Small improvements to the tuya ts0021 quirk

tests/test_tuya.py

@@ -50,6 +50,7 @@
 ZCL_TUYA_BUTTON_2_SINGLE_PRESS = b"\tN\x06\x01\x1f\x02\x04\x00\x01\x00"
 ZCL_TUYA_BUTTON_2_DOUBLE_PRESS = b"\tj\x06\x03\x10\x02\x04\x00\x01\x01"
 ZCL_TUYA_BUTTON_2_LONG_PRESS = b"\tl\x06\x03\x12\x02\x04\x00\x01\x02"
+ZCL_TUYA_BATTERY_REPORT_100_PCT = b"\tv\x06\x00\xed\n\x02\x00\x04\x00\x00\x00d"
 ZCL_TUYA_SWITCH_ON = b"\tQ\x02\x006\x01\x01\x00\x01\x01"
 ZCL_TUYA_SWITCH_OFF = b"\tQ\x02\x006\x01\x01\x00\x01\x00"
 ZCL_TUYA_ATTRIBUTE_617_TO_179 = b"\tp\x02\x00\x02i\x02\x00\x04\x00\x00\x00\xb3"
@@ -116,65 +117,126 @@ async def test_singleswitch_state_report(zigpy_device_from_quirk, quirk):
 
 
 @pytest.mark.parametrize(
-    "quirk,raw_event,expected_attr_name,expected_attr_value",
+    "quirk,raw_event,button_number,press_type,expected_attr_value",
     (
         (
             zhaquirks.tuya.ts0021.TS0021,
             ZCL_TUYA_BUTTON_1_SINGLE_PRESS,
-            "btn_1_pressed",
+            1,
+            "single",
             0x00,
         ),
         (
             zhaquirks.tuya.ts0021.TS0021,
             ZCL_TUYA_BUTTON_1_DOUBLE_PRESS,
-            "btn_1_pressed",
+            1,
+            "double",
             0x01,
         ),
         (
             zhaquirks.tuya.ts0021.TS0021,
             ZCL_TUYA_BUTTON_1_LONG_PRESS,
-            "btn_1_pressed",
+            1,
+            "long",
             0x02,
         ),
         (
             zhaquirks.tuya.ts0021.TS0021,
             ZCL_TUYA_BUTTON_2_SINGLE_PRESS,
-            "btn_2_pressed",
+            2,
+            "single",
             0x00,
         ),
         (
             zhaquirks.tuya.ts0021.TS0021,
             ZCL_TUYA_BUTTON_2_DOUBLE_PRESS,
-            "btn_2_pressed",
+            2,
+            "double",
             0x01,
         ),
         (
             zhaquirks.tuya.ts0021.TS0021,
             ZCL_TUYA_BUTTON_2_LONG_PRESS,
-            "btn_2_pressed",
+            2,
+            "long",
             0x02,
         ),
     ),
 )
-async def test_ts0021_switch(
-    zigpy_device_from_quirk, quirk, raw_event, expected_attr_name, expected_attr_value
+async def test_ts0021_switch_button_press(
+    zigpy_device_from_quirk,
+    quirk,
+    raw_event,
+    button_number,
+    press_type,
+    expected_attr_value,
 ):
     """Test tuya TS0021 2-gang switch."""
 
     device = zigpy_device_from_quirk(quirk)
 
     tuya_cluster = device.endpoints[1].tuya_manufacturer
     switch_listener = ClusterListener(tuya_cluster)
+    listener = mock.MagicMock()
+    tuya_cluster.add_listener(listener)
 
     hdr, args = tuya_cluster.deserialize(raw_event)
     tuya_cluster.handle_message(hdr, args)
 
     assert len(switch_listener.cluster_commands) == 1
     assert len(switch_listener.attribute_updates) == 1
 
-    assert switch_listener.attribute_updates[0][0] == expected_attr_name
+    assert switch_listener.attribute_updates[0][0] == f"btn_{button_number}_pressed"
     assert switch_listener.attribute_updates[0][1] == expected_attr_value
 
+    # We fire two events:
+    # - One for the attribute update;
+    # - One for the button press.
+    assert listener.zha_send_event.mock_calls == [
+        mock.call(
+            "attribute_updated",
+            {
+                "attribute_id": f"btn_{button_number}_pressed",
+                "attribute_name": "Unknown",
+                "value": expected_attr_value,
+            },
+        ),
+        mock.call(
+            f"button_{button_number}_{press_type}_press",
+            {
+                "button": button_number,
+                "press_type": press_type,
+            },
+        ),
+    ]
+
+
+async def test_ts0021_switch_battery_report(zigpy_device_from_quirk):
+    """Test tuya TS0021 2-gang switch."""
+
+    device = zigpy_device_from_quirk(zhaquirks.tuya.ts0021.TS0021)
+
+    tuya_cluster = device.endpoints[1].tuya_manufacturer
+    power_cluster = device.endpoints[1].power
+
+    switch_listener = ClusterListener(tuya_cluster)
+    power_listener = ClusterListener(power_cluster)
+
+    tuya_cluster.add_listener(mock.MagicMock())
+    power_cluster.add_listener(mock.MagicMock())
+
+    hdr, args = tuya_cluster.deserialize(ZCL_TUYA_BATTERY_REPORT_100_PCT)
+    tuya_cluster.handle_message(hdr, args)
+
+    # Custom manufacturer cluster handled 1 message.
+    assert len(switch_listener.cluster_commands) == 1
+
+    # We do not update the attribute internally, but instead pipe it to the PowerConfiguration cluster.
+    assert len(switch_listener.attribute_updates) == 0
+
+    # PowerConfiguration.battery_percentage_remaining.id, 100% in 0.5% increments == 200.
+    assert power_listener.attribute_updates == [(0x0021, 200)]
+
 
 @pytest.mark.parametrize("quirk", (zhaquirks.tuya.ts0601_switch.TuyaDoubleSwitchTO,))
 async def test_doubleswitch_state_report(zigpy_device_from_quirk, quirk):

zhaquirks/tuya/ts0021.py

@@ -5,7 +5,7 @@
 from zigpy.zcl.clusters.general import Basic, Ota, PowerConfiguration, Time
 from zigpy.zcl.clusters.security import IasZone
 
-from zhaquirks import EventableCluster
+from zhaquirks import Bus, EventableCluster
 from zhaquirks.const import (
     ARGS,
     ATTRIBUTE_ID,
@@ -24,7 +24,13 @@
     SHORT_PRESS,
     VALUE,
 )
-from zhaquirks.tuya import TUYA_CLUSTER_ID, DPToAttributeMapping, TuyaNewManufCluster
+from zhaquirks.tuya import (
+    TUYA_CLUSTER_ID,
+    DPToAttributeMapping,
+    TuyaDatapointData,
+    TuyaNewManufCluster,
+    TuyaPowerConfigurationCluster2AAA,
+)
 
 BTN_1 = "Button 1"
 BTN_2 = "Button 2"
@@ -36,6 +42,8 @@
 class TuyaCustomCluster(TuyaNewManufCluster, EventableCluster):
     """Tuya Custom Cluster for mapping data points to attributes."""
 
+    PRESS_TYPE = {0: "single", 1: "double", 2: "long"}
+
     dp_to_attribute: dict[int, DPToAttributeMapping] = {
         1: DPToAttributeMapping(
             TuyaNewManufCluster.ep_attribute,
@@ -50,12 +58,53 @@ class TuyaCustomCluster(TuyaNewManufCluster, EventableCluster):
     data_point_handlers = {
         1: "_dp_2_attr_update",
         2: "_dp_2_attr_update",
+        10: "_battery_pct_attr_update",
     }
 
+    def _dp_2_attr_update(self, datapoint: TuyaDatapointData) -> None:
+        super()._dp_2_attr_update(datapoint)
+        button_n = datapoint.dp
+        press_type = self.PRESS_TYPE.get(datapoint.data.payload, "unknown")
+        action = f"button_{button_n}_{press_type}_press"
+        self.listener_event(
+            "zha_send_event",
+            action,
+            {
+                "button": button_n,
+                "press_type": press_type,
+            },
+        )
+
+    def _battery_pct_attr_update(self, datapoint: TuyaDatapointData) -> None:
+        self.endpoint.device.battery_pct_bus.listener_event(
+            "battery_percentage_reported", datapoint.data.payload
+        )
+
+
+class TuyaCustomPowerCluster(TuyaPowerConfigurationCluster2AAA):
+    """Tuya Custom PowerCluster. This cluster is used to report battery percentage."""
+
+    def __init__(self, *args, **kwargs):
+        """Init cluster."""
+        super().__init__(*args, **kwargs)
+        self.endpoint.device.battery_pct_bus.add_listener(self)
+
+    def battery_percentage_reported(self, value: int) -> None:
+        """Handle battery percentage reported."""
+        # Reports battery percentage in 0.5% increments; i.e. 2 x the actual percentage.
+        self._update_attribute(
+            self.AttributeDefs.battery_percentage_remaining.id, 2 * value
+        )
+
 
 class TS0021(CustomDevice):
     """Tuya TS0021 2-button switch device."""
 
+    def __init__(self, *args, **kwargs):
+        """Init device."""
+        self.battery_pct_bus = Bus()
+        super().__init__(*args, **kwargs)
+
     signature = {
         # SizePrefixedSimpleDescriptor(endpoint=1, profile=260, device_type=1026,
         # device_version=1,
@@ -87,8 +136,7 @@ class TS0021(CustomDevice):
                 DEVICE_TYPE: zha.DeviceType.IAS_ZONE,
                 INPUT_CLUSTERS: [
                     Basic.cluster_id,
-                    PowerConfiguration.cluster_id,
-                    IasZone.cluster_id,
+                    TuyaCustomPowerCluster,
                     TuyaCustomCluster,
                 ],
                 OUTPUT_CLUSTERS: [