update tests

zigpy · Jan 23, 2025 · a97f6d9 · a97f6d9
1 parent 6405d78
commit a97f6d9
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/tests/test_tuya_builder.py b/tests/test_tuya_builder.py
@@ -8,10 +8,11 @@
 from zigpy.quirks.v2 import CustomDeviceV2
 import zigpy.types as t
 from zigpy.zcl import foundation
-from zigpy.zcl.clusters.general import Basic, BatterySize
+from zigpy.zcl.clusters.general import Basic
 
 from tests.common import ClusterListener, MockDatetime, wait_for_zigpy_tasks
 import zhaquirks
+from zhaquirks.const import BatterySize
 from zhaquirks.tuya import (
     TUYA_QUERY_DATA,
     TUYA_SET_TIME,
@@ -93,6 +94,9 @@ async def test_convenience_methods(device_mock, method_name, attr_name, exp_clas
         (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.CR2450, 1, None, BatterySize.CR2450, 1, 30),
+        (None, BatterySize.CR2032, 1, None, BatterySize.CR2032, 1, 30),
+        (None, BatterySize.CR1632, 1, None, BatterySize.CR1632, 1, 30),
         (None, BatterySize.AA, None, None, BatterySize.AA, None, None),
         (None, None, None, None, None, None, None),
     ],

diff --git a/tests/test_xiaomi.py b/tests/test_xiaomi.py
@@ -47,6 +47,7 @@
     OUTPUT_CLUSTERS,
     PROFILE_ID,
     ZONE_STATUS_CHANGE_COMMAND,
+    BatterySize,
 )
 from zhaquirks.xiaomi import (
     LUMI,
@@ -384,10 +385,10 @@ async def test_mija_battery(zigpy_device_from_quirk, voltage, bpr):
 @pytest.mark.parametrize(
     "quirk, batt_size",
     (
-        (zhaquirks.xiaomi.aqara.vibration_aq1.VibrationAQ1, 0x0A),
-        (zhaquirks.xiaomi.mija.motion.Motion, 0x09),
-        (zhaquirks.xiaomi.mija.sensor_switch.MijaButton, 0x0A),
-        (zhaquirks.xiaomi.mija.sensor_magnet.Magnet, 0x0B),
+        (zhaquirks.xiaomi.aqara.vibration_aq1.VibrationAQ1, BatterySize.CR2032),
+        (zhaquirks.xiaomi.mija.motion.Motion, BatterySize.CR2450),
+        (zhaquirks.xiaomi.mija.sensor_switch.MijaButton, BatterySize.CR2032),
+        (zhaquirks.xiaomi.mija.sensor_magnet.Magnet, BatterySize.CR1632),
     ),
 )
 async def test_xiaomi_batt_size(zigpy_device_from_quirk, quirk, batt_size):