Repository: zigpy/zha-device-handlers
Pull request #3728 (2025-01-19 00:08:46)

Enhance TuyaQuirkBuilder battery method
## Proposed change
<!--
  Explain your proposed change below.
-->

As we add v2 Tuya quirks, we also keep adding PowerConfiguration clusters. This allows calling `tuya_battery` and specifying the battery type, quantity, and voltage. 

If voltage is omitted, it will be looked up from a dict of known voltages.

Maintains backwards compatibility by allowing a PowerConfiguration cluster to be specified.

```python
.tuya_battery(dp_id=2, battery_type=BatterySize.AA, battery_qty=4)
```

## Additional information
<!--
  Please include any additional information that is important to this PR.
  For example, if this PR is a potentially breaking change, mention that here.
  If this PR requires other PRs to be merged in HA Core or other projects, mention that.
  Lastly, if this PR fixes a specific issue, please include "Fixes #xxxx".
-->


## Checklist
<!--
  Put an 'x' in all boxes that apply.
  Note: You do not need to tick all boxes before creating a PR.
-->

- [x] The changes are tested and work correctly
- [x] `pre-commit` checks pass / the code has been formatted using Black
- [x] Tests have been added to verify that the new code works


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/zigpy/zha-device-handlers/pull/3728?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zigpy) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 89.88%. Comparing base [(`e61ecbb`)](https://app.codecov.io/gh/zigpy/zha-device-handlers/commit/e61ecbb4f64291c663fd19d23f95f7db5d33e1f0?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zigpy) to head [(`6f60ffe`)](https://app.codecov.io/gh/zigpy/zha-device-handlers/commit/6f60ffe16e3f310f8d4053cd100dba8e5dbd731f?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zigpy).

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##              dev    #3728      +/-   ##
==========================================
+ Coverage   89.87%   89.88%   +0.01%     
==========================================
  Files         320      320              
  Lines       10381    10392      +11     
==========================================
+ Hits         9330     9341      +11     
  Misses       1051     1051              
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/zigpy/zha-device-handlers/pull/3728?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zigpy).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zigpy).


Patch:
diff --git a/tuya.md b/tuya.md
index 02a79ec..f77458f 100644
--- a/tuya.md
+++ b/tuya.md
@@ -56,12 +56,18 @@ TuyaQuirkBuilder is a subclass of QuirkBuilder, retaining all of the v2 QuirkBui
 
 These methods allow exposing the most common Tuya clusters. These methods were added as part of the quirk building process and it is likely that there are other convenience methods that should be created. If you find that you are repeating the `.tuya_dp` and `.adds` formula, please PR or suggest additional methods.
 
-#### tuya_battery(dp_id: int, power_cfg: PowerConfiguration = TuyaPowerConfigurationCluster2AAA, scale: float = 2)
+#### def tuya_battery(
+####        dp_id: int,
+####        power_cfg: PowerConfiguration | None = None,
+####        battery_type: BatterySize | None = BatterySize.AA,
+####        battery_qty: int | None = 2,
+####        battery_voltage: int | None = None,
+####        scale: float = 2,
 
 Adds a battery power cluster.
 
 ```python
-.tuya_battery(dp_id=2, power_config=TuyaPowerConfigurationCluster4AAA)
+.tuya_battery(dp_id=2, battery_type=BatterySize.AA, battery_qty=4)
 ```
 
 #### tuya_metering(dp_id: int, metering_cfg: TuyaLocalCluster = TuyaValveWaterConsumed)
diff --git a/zhaquirks/tuya/builder/__init__.py b/zhaquirks/tuya/builder/__init__.py
index 638752f..64780d2 100644
--- a/zhaquirks/tuya/builder/__init__.py
+++ b/zhaquirks/tuya/builder/__init__.py
@@ -13,6 +13,7 @@ from zigpy.quirks.v2.homeassistant.number import NumberDeviceClass
 from zigpy.quirks.v2.homeassistant.sensor import SensorDeviceClass, SensorStateClass
 import zigpy.types as t
 from zigpy.zcl import foundation
+from zigpy.zcl.clusters.general import BatterySize
 from zigpy.zcl.clusters.measurement import (
     PM25,
     CarbonDioxideConcentration,
@@ -30,13 +31,28 @@ from zhaquirks.tuya import (
     BaseEnchantedDevice,
     PowerConfiguration,
     TuyaLocalCluster,
-    TuyaPowerConfigurationCluster2AAA,
+    TuyaPowerConfigurationCluster,
 )
 from zhaquirks.tuya.mcu import DPToAttributeMapping, TuyaMCUCluster, TuyaOnOffNM
 
 MOL_VOL_AIR_NTP = 0.2445  # molar volume of air at NTP in cL/mol
 
 
+BATTERY_VOLTAGES = {
+    BatterySize.No_battery: None,
+    BatterySize.Built_in: None,
+    BatterySize.Other: None,
+    BatterySize.AA: 15,
+    BatterySize.AAA: 15,
+    BatterySize.C: 15,
+    BatterySize.D: 15,
+    BatterySize.AA: 15,
+    BatterySize.CR2: 30,
+    BatterySize.CR123A: 30,
+    BatterySize.Unknown: None,
+}
+
+
 class TuyaCO2Concetration(CarbonDioxideConcentration, TuyaLocalCluster):
     """Tuya Carbon Dioxide concentration measurement."""
 
@@ -148,22 +164,52 @@ class TuyaQuirkBuilder(QuirkBuilder):
         self.new_attributes: set[foundation.ZCLAttributeDef] = set()
         super().__init__(manufacturer, model, registry)
 
-    def tuya_battery(
+    def _tuya_battery(
         self,
         dp_id: int,
-        power_cfg: PowerConfiguration = TuyaPowerConfigurationCluster2AAA,
-        scale: float = 2,
+        power_cfg: PowerConfiguration,
+        scale: float,
     ) -> QuirkBuilder:
         """Add a Tuya Battery Power Configuration."""
         self.tuya_dp(
             dp_id,
             power_cfg.ep_attribute,
-            "battery_percentage_remaining",
+            PowerConfiguration.AttributeDefs.battery_percentage_remaining.name,
             converter=lambda x: x * scale,
         )
         self.adds(power_cfg)
         return self
 
+    def tuya_battery(
+        self,
+        dp_id: int,
+        power_cfg: PowerConfiguration | None = None,
+        battery_type: BatterySize | None = BatterySize.AA,
+        battery_qty: int | None = 2,
+        battery_voltage: int | None = None,
+        scale: float = 2,
+    ) -> QuirkBuilder:
+        """Add a Tuya Battery Power Configuration."""
+
+        if power_cfg:
+            return self._tuya_battery(dp_id=dp_id, power_cfg=power_cfg, scale=scale)
+
+        if not battery_voltage and (battery_type and battery_qty):
+            battery_voltage = BATTERY_VOLTAGES.get(battery_type)
+
+        class TuyaPowerConfigurationClusterBattery(TuyaPowerConfigurationCluster):
+            """PowerConfiguration cluster for Tuya devices."""
+
+            _CONSTANT_ATTRIBUTES = {
+                PowerConfiguration.AttributeDefs.battery_size.id: battery_type,
+                PowerConfiguration.AttributeDefs.battery_rated_voltage.id: battery_voltage,
+                PowerConfiguration.AttributeDefs.battery_quantity.id: battery_qty,
+            }
+
+        return self._tuya_battery(
+            dp_id=dp_id, power_cfg=TuyaPowerConfigurationClusterBattery, scale=scale
+        )
+
     def tuya_contact(self, dp_id: int):
         """Add a Tuya IAS contact sensor."""
         self.tuya_ias(
diff --git a/zhaquirks/tuya/ts0601_sensor.py b/zhaquirks/tuya/ts0601_sensor.py
index e8a0674..e68cff5 100644
--- a/zhaquirks/tuya/ts0601_sensor.py
+++ b/zhaquirks/tuya/ts0601_sensor.py
@@ -3,11 +3,8 @@
 from zigpy.quirks.v2.homeassistant.sensor import SensorDeviceClass, SensorStateClass
 import zigpy.types as t
 
-from zhaquirks.tuya.builder import (
-    TuyaPowerConfigurationCluster2AAA,
-    TuyaQuirkBuilder,
-    TuyaTemperatureMeasurement,
-)
+from zhaquirks.tuya import TuyaPowerConfigurationCluster2AAA
+from zhaquirks.tuya.builder import TuyaQuirkBuilder, TuyaTemperatureMeasurement
 
 (
     TuyaQuirkBuilder("_TZE200_bjawzodf", "TS0601")
diff --git a/zhaquirks/tuya/ts0601_valve.py b/zhaquirks/tuya/ts0601_valve.py
index 2434435..26f5be2 100644
--- a/zhaquirks/tuya/ts0601_valve.py
+++ b/zhaquirks/tuya/ts0601_valve.py
@@ -6,12 +6,9 @@ from zigpy.quirks.v2 import EntityPlatform, EntityType
 from zigpy.quirks.v2.homeassistant import UnitOfTime
 from zigpy.quirks.v2.homeassistant.sensor import SensorDeviceClass, SensorStateClass
 import zigpy.types as t
+from zigpy.zcl.clusters.general import BatterySize
 
-from zhaquirks.tuya import (
-    TUYA_CLUSTER_ID,
-    TuyaPowerConfigurationCluster2AA,
-    TuyaPowerConfigurationCluster4AA,
-)
+from zhaquirks.tuya import TUYA_CLUSTER_ID
 from zhaquirks.tuya.builder import TuyaQuirkBuilder
 from zhaquirks.tuya.mcu import TuyaMCUCluster
 
@@ -45,7 +42,7 @@ class TuyaValveTimerState(t.enum8):
         attribute_name="dp_6",
         type=t.uint32_t,
     )
-    .tuya_battery(dp_id=7, power_cfg=TuyaPowerConfigurationCluster4AA)
+    .tuya_battery(dp_id=7, battery_type=BatterySize.AA, battery_qty=4)
     .tuya_enum(
         dp_id=10,
         attribute_name="weather_delay",
@@ -134,7 +131,7 @@ class ParksideTuyaValveManufCluster(TuyaMCUCluster):
         translation_key="time_left",
         fallback_name="Time left",
     )
-    .tuya_battery(dp_id=11, power_cfg=TuyaPowerConfigurationCluster4AA)
+    .tuya_battery(dp_id=11, battery_type=BatterySize.AA, battery_qty=4)
     .tuya_switch(
         dp_id=108,
         attribute_name="frost_lock",
@@ -202,7 +199,7 @@ def giex_string_to_ts(v: str) -> int | None:
 
 gx02_base_quirk = (
     TuyaQuirkBuilder()
-    .tuya_battery(dp_id=108, power_cfg=TuyaPowerConfigurationCluster4AA)
+    .tuya_battery(dp_id=108, battery_type=BatterySize.AA, battery_qty=4)
     .tuya_metering(dp_id=111)
     .tuya_onoff(dp_id=2)
     .tuya_number(
@@ -337,7 +334,7 @@ class GiexIrrigationStatus(t.enum8):
 (
     TuyaQuirkBuilder("_TZE284_8zizsafo", "TS0601")  # Giex GX04
     .applies_to("_TZE284_eaet5qt5", "TS0601")  # Insoma SGW08W
-    .tuya_battery(dp_id=59, power_cfg=TuyaPowerConfigurationCluster4AA)
+    .tuya_battery(dp_id=59, battery_type=BatterySize.AA, battery_qty=4)
     .tuya_switch(
         dp_id=1,
         attribute_name="valve_on_off_1",
@@ -424,7 +421,7 @@ class GiexIrrigationStatus(t.enum8):
 (
     TuyaQuirkBuilder("_TZE200_2wg5qrjy", "TS0601")
     .tuya_onoff(dp_id=1)
-    .tuya_battery(dp_id=7, power_cfg=TuyaPowerConfigurationCluster2AA)
+    .tuya_battery(dp_id=7, battery_type=BatterySize.AA, battery_qty=2)
     # Water consumed (value comes in deciliters - convert it to liters)
     .tuya_metering(dp_id=5, scale=0.1)
     # Timer time left/remaining (raw value in seconds)
diff --git a/zhaquirks/tuya/tuya_siren.py b/zhaquirks/tuya/tuya_siren.py
index 132d9a9..ed9cc3f 100644
--- a/zhaquirks/tuya/tuya_siren.py
+++ b/zhaquirks/tuya/tuya_siren.py
@@ -4,8 +4,8 @@ from zigpy.quirks.v2 import EntityPlatform, EntityType
 from zigpy.quirks.v2.homeassistant import UnitOfTime
 from zigpy.quirks.v2.homeassistant.binary_sensor import BinarySensorDeviceClass
 import zigpy.types as t
+from zigpy.zcl.clusters.general import BatterySize
 
-from zhaquirks.tuya import TuyaPowerConfigurationClusterOther
 from zhaquirks.tuya.builder import TuyaQuirkBuilder
 
 
@@ -62,7 +62,9 @@ class TuyaSirenRingtone(t.enum8):
         translation_key="siren_on",
         fallback_name="Siren on",
     )
-    .tuya_battery(dp_id=15, power_cfg=TuyaPowerConfigurationClusterOther)
+    .tuya_battery(
+        dp_id=15, battery_type=BatterySize.Other, battery_qty=1, battery_voltage=30
+    )
     .tuya_binary_sensor(
         dp_id=20,
         attribute_name="tamper_state",
diff --git a/zhaquirks/tuya/tuya_smoke.py b/zhaquirks/tuya/tuya_smoke.py
index 39beaa1..c4e9da4 100644
--- a/zhaquirks/tuya/tuya_smoke.py
+++ b/zhaquirks/tuya/tuya_smoke.py
@@ -9,8 +9,8 @@ from zigpy.zcl.clusters.security import IasZone
 from zigpy.zcl.foundation import BaseAttributeDefs, ZCLAttributeDef
 
 from zhaquirks import LocalDataCluster
-from zhaquirks.tuya import TuyaManufClusterAttributes
-from zhaquirks.tuya.builder import TuyaPowerConfigurationCluster2AAA, TuyaQuirkBuilder
+from zhaquirks.tuya import TuyaManufClusterAttributes, TuyaPowerConfigurationCluster2AAA
+from zhaquirks.tuya.builder import TuyaQuirkBuilder
 
 
 class TuyaIasZone(LocalDataCluster, IasZone):


Test patch:
diff --git a/tests/test_tuya_builder.py b/tests/test_tuya_builder.py
index 902faf6..99031bf 100644
--- a/tests/test_tuya_builder.py
+++ b/tests/test_tuya_builder.py
@@ -7,11 +7,15 @@ from zigpy.quirks.registry import DeviceRegistry
 from zigpy.quirks.v2 import CustomDeviceV2
 import zigpy.types as t
 from zigpy.zcl import foundation
-from zigpy.zcl.clusters.general import Basic
+from zigpy.zcl.clusters.general import Basic, BatterySize
 
 from tests.common import ClusterListener, wait_for_zigpy_tasks
 import zhaquirks
-from zhaquirks.tuya import TUYA_QUERY_DATA
+from zhaquirks.tuya import (
+    TUYA_QUERY_DATA,
+    TuyaPowerConfigurationCluster,
+    TuyaPowerConfigurationCluster2AAA,
+)
 from zhaquirks.tuya.builder import (
     TuyaAirQualityVOC,
     TuyaCO2Concetration,
@@ -19,7 +23,6 @@ from zhaquirks.tuya.builder import (
     TuyaIasContact,
     TuyaIasFire,
     TuyaPM25Concetration,
-    TuyaPowerConfigurationCluster2AAA,
     TuyaQuirkBuilder,
     TuyaRelativeHumidity,
     TuyaSoilMoisture,
@@ -34,7 +37,7 @@ zhaquirks.setup()
 @pytest.mark.parametrize(
     "method_name,attr_name,exp_class",
     [
-        ("tuya_battery", "power", TuyaPowerConfigurationCluster2AAA),
+        ("tuya_battery", "power", TuyaPowerConfigurationCluster),
         ("tuya_metering", "smartenergy_metering", TuyaValveWaterConsumed),
         ("tuya_onoff", "on_off", TuyaOnOffNM),
         ("tuya_soil_moisture", "soil_moisture", TuyaSoilMoisture),
@@ -74,6 +77,56 @@ async def test_convenience_methods(device_mock, method_name, attr_name, exp_clas
     assert isinstance(ep_attr, exp_class)
 
 
+@pytest.mark.parametrize(
+    "power_cfg,battery_type,battery_qty,battery_voltage,"
+    "expected_size,expected_qty,expected_voltage",
+    [
+        (TuyaPowerConfigurationCluster2AAA, None, None, None, BatterySize.AAA, 2, 15),
+        (None, BatterySize.CR123A, 1, 60, BatterySize.CR123A, 1, 60),
+        (None, BatterySize.CR123A, 1, None, BatterySize.CR123A, 1, 30),
+        (None, BatterySize.AA, None, None, BatterySize.AA, None, None),
+        (None, None, None, None, None, None, None),
+    ],
+)
+async def test_battery_methods(
+    device_mock,
+    power_cfg,
+    battery_type,
+    battery_qty,
+    battery_voltage,
+    expected_size,
+    expected_qty,
+    expected_voltage,
+):
+    """Test the battery convenience method."""
+
+    registry = DeviceRegistry()
+
+    (
+        TuyaQuirkBuilder(device_mock.manufacturer, device_mock.model, registry=registry)
+        .tuya_battery(
+            dp_id=1,
+            power_cfg=power_cfg,
+            battery_type=battery_type,
+            battery_qty=battery_qty,
+            battery_voltage=battery_voltage,
+        )
+        .tuya_onoff(dp_id=3)
+        .skip_configuration()
+        .add_to_registry()
+    )
+
+    quirked = registry.get_device(device_mock)
+    ep = quirked.endpoints[1]
+
+    assert ep.power is not None
+    assert isinstance(ep.power, power_cfg or TuyaPowerConfigurationCluster)
+
+    assert ep.power.get("battery_size") == expected_size
+    assert ep.power.get("battery_quantity") == expected_qty
+    assert ep.power.get("battery_rated_voltage") == expected_voltage
+
+
 async def test_tuya_quirkbuilder(device_mock):
     """Test adding a v2 Tuya Quirk to the registry and getting back a quirked device."""