Init quirk

Author: prairiesnpr

tests/conftest.py

@@ -5,6 +5,7 @@
 import pytest
 import zigpy.application
 import zigpy.device
+from zigpy.device import Device
 import zigpy.quirks
 import zigpy.types
 from zigpy.zcl import foundation
@@ -21,6 +22,8 @@
     PROFILE_ID,
 )
 
+from .async_mock import sentinel
+
 
 class MockApp(zigpy.application.ControllerApplication):
     """App Controller."""
@@ -204,6 +207,25 @@ def _dev(
     return _dev
 
 
+@pytest.fixture(name="device_mock")
+def real_device(MockAppController):
+    """Device fixture with a single endpoint."""
+    ieee = sentinel.ieee
+    nwk = 0x2233
+    device = Device(MockAppController, ieee, nwk)
+
+    device.add_endpoint(1)
+    device[1].profile_id = 0x0104
+    device[1].device_type = 0x0051
+    device.model = "model"
+    device.manufacturer = "manufacturer"
+    device[1].add_input_cluster(0x0000)
+    device[1].add_input_cluster(0xEF00)
+    device[1].add_output_cluster(0x000A)
+    device[1].add_output_cluster(0x0019)
+    return device
+
+
 @pytest.fixture
 def assert_signature_matches_quirk():
     """Return a function that can be used to check if a given quirk matches a signature."""

tests/test_quirks.py

@@ -14,7 +14,7 @@
 import zigpy.endpoint
 import zigpy.profiles
 import zigpy.quirks as zq
-from zigpy.quirks import CustomDevice
+from zigpy.quirks import CustomDevice, DeviceRegistry
 from zigpy.quirks.v2 import QuirkBuilder
 import zigpy.types
 from zigpy.zcl import foundation
@@ -845,8 +845,9 @@ def check_for_duplicate_cluster_ids(clusters) -> None:
         check_for_duplicate_cluster_ids(ep_data.get(OUTPUT_CLUSTERS, []))
 
 
-async def test_local_data_cluster(zigpy_device_from_v2_quirk) -> None:
+async def test_local_data_cluster(device_mock) -> None:
     """Ensure reading attributes from a LocalDataCluster works as expected."""
+    registry = DeviceRegistry()
 
     class TestLocalCluster(zhaquirks.LocalDataCluster):
         """Test cluster."""
@@ -856,11 +857,11 @@ class TestLocalCluster(zhaquirks.LocalDataCluster):
         _VALID_ATTRIBUTES = [2]
 
     (
-        QuirkBuilder("manufacturer-local-test", "model")
+        QuirkBuilder(device_mock.manufacturer, device_mock.model, registry=registry)
         .adds(TestLocalCluster)
         .add_to_registry()
     )
-    device = zigpy_device_from_v2_quirk("manufacturer-local-test", "model")
+    device = registry.get_device(device_mock)
     assert isinstance(device.endpoints[1].in_clusters[0x1234], TestLocalCluster)
 
     # reading invalid attribute return unsupported attribute

tests/test_tuya_builder.py

@@ -3,7 +3,6 @@
 from unittest import mock
 
 import pytest
-from zigpy.device import Device
 from zigpy.quirks.registry import DeviceRegistry
 from zigpy.quirks.v2 import CustomDeviceV2
 import zigpy.types as t
@@ -25,30 +24,9 @@
 )
 from zhaquirks.tuya.mcu import TuyaMCUCluster, TuyaOnOffNM
 
-from .async_mock import sentinel
-
 zhaquirks.setup()
 
 
-@pytest.fixture(name="device_mock")
-def real_device(MockAppController):
-    """Device fixture with a single endpoint."""
-    ieee = sentinel.ieee
-    nwk = 0x2233
-    device = Device(MockAppController, ieee, nwk)
-
-    device.add_endpoint(1)
-    device[1].profile_id = 0x0104
-    device[1].device_type = 0x0051
-    device.model = "model"
-    device.manufacturer = "manufacturer"
-    device[1].add_input_cluster(0x0000)
-    device[1].add_input_cluster(0xEF00)
-    device[1].add_output_cluster(0x000A)
-    device[1].add_output_cluster(0x0019)
-    return device
-
-
 @pytest.mark.parametrize(
     "method_name,attr_name,exp_class",
     [
@@ -95,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)
@@ -133,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
@@ -150,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")