Init quirk

Author: prairiesnpr

tests/test_tuya_builder.py

@@ -73,6 +73,9 @@ class TestEnum(t.enum8):
         A = 0x00
         B = 0x01
 
+    class ModTuyaMCUCluster(TuyaMCUCluster):
+        """Modified Cluster."""
+
     entry = (
         TuyaQuirkBuilder(device_mock.manufacturer, device_mock.model, registry=registry)
         .tuya_battery(dp_id=1)
@@ -111,7 +114,7 @@ class TestEnum(t.enum8):
             fallback_name="Test enum",
         )
         .skip_configuration()
-        .add_to_registry()
+        .add_to_registry(mcu_cluster=ModTuyaMCUCluster)
     )
 
     # coverage for overridden __eq__ method
@@ -128,6 +131,7 @@ class TestEnum(t.enum8):
     assert isinstance(ep.basic, Basic)
 
     assert ep.tuya_manufacturer is not None
+    assert isinstance(ep.tuya_manufacturer, ModTuyaMCUCluster)
     assert isinstance(ep.tuya_manufacturer, TuyaMCUCluster)
 
     tuya_cluster = ep.tuya_manufacturer

zhaquirks/tuya/builder/__init__.py

@@ -538,16 +538,16 @@ class EnchantedDeviceV2(CustomDeviceV2, BaseEnchantedDevice):
 
         return self
 
-    def add_to_registry(self) -> QuirksV2RegistryEntry:
+    def add_to_registry(self, mcu_cluster=TuyaMCUCluster) -> QuirksV2RegistryEntry:
         """Build the quirks v2 registry entry."""
 
-        class NewAttributeDefs(TuyaMCUCluster.AttributeDefs):
+        class NewAttributeDefs(mcu_cluster.AttributeDefs):
             """Attribute Definitions."""
 
         for attr in self.new_attributes:
             setattr(NewAttributeDefs, attr.name, attr)
 
-        class TuyaReplacementCluster(TuyaMCUCluster):
+        class TuyaReplacementCluster(mcu_cluster):
             """Replacement Tuya Cluster."""
 
             data_point_handlers: dict[int, str]

zhaquirks/tuya/ts0601_sensor.py

@@ -2,8 +2,15 @@
 
 from zigpy.quirks.v2.homeassistant.sensor import SensorDeviceClass, SensorStateClass
 import zigpy.types as t
+from zigpy.zcl import foundation
 
-from zhaquirks.tuya.builder import TuyaPowerConfigurationCluster2AAA, TuyaQuirkBuilder
+from zhaquirks.tuya import (
+    TUYA_MCU_VERSION_REQ,
+    TUYA_QUERY_DATA,
+    TuyaPowerConfigurationCluster2AAA,
+)
+from zhaquirks.tuya.builder import TuyaQuirkBuilder
+from zhaquirks.tuya.mcu import TuyaMCUCluster
 
 (
     TuyaQuirkBuilder("_TZE200_bjawzodf", "TS0601")
@@ -33,6 +40,49 @@
 )
 
 
+class RespondingTuyaMCUCluster(TuyaMCUCluster):
+    """Tuya Manufacturer Cluster with mcu version response."""
+
+    def handle_mcu_version_response(
+        self, payload: TuyaMCUCluster.MCUVersion
+    ) -> foundation.Status:
+        """Handle MCU version response."""
+
+        self.create_catching_task(
+            super().command(TUYA_MCU_VERSION_REQ, 2, expect_reply=False)
+        )
+        self.create_catching_task(super().command(TUYA_QUERY_DATA, expect_reply=False))
+        super().handle_mcu_version_response(payload)
+
+    server_commands = TuyaMCUCluster.server_commands.copy()
+    server_commands.update(
+        {
+            TUYA_MCU_VERSION_REQ: foundation.ZCLCommandDef(
+                "mcu_version_req",
+                {"data": t.uint16_t},
+                True,
+                is_manufacturer_specific=False,
+            ),
+        }
+    )
+
+
+(
+    TuyaQuirkBuilder("_TZE204_upagmta9", "TS0601")
+    .tuya_temperature(dp_id=1, scale=10)
+    .tuya_humidity(dp_id=2)
+    .tuya_dp(
+        dp_id=3,
+        ep_attribute=TuyaPowerConfigurationCluster2AAA.ep_attribute,
+        attribute_name="battery_percentage_remaining",
+        converter=lambda x: {0: 50, 1: 100, 2: 200}[x],
+    )
+    .adds(TuyaPowerConfigurationCluster2AAA)
+    .skip_configuration()
+    .add_to_registry(mcu_cluster=RespondingTuyaMCUCluster)
+)
+
+
 (
     TuyaQuirkBuilder("_TZE200_yjjdcqsq", "TS0601")
     .applies_to("_TZE200_9yapgbuv", "TS0601")