Initial tests

Author: prairiesnpr

tests/async_mock.py

@@ -0,0 +1,43 @@
+"""Mock utilities that are async aware."""
+
+from unittest.mock import *  # noqa: F401, F403
+
+
+class _IntSentinelObject(int):
+    """Sentinel-like object that is also an integer subclass.
+
+    Allows sentinels to be used
+    in loggers that perform int-specific string formatting.
+    """
+
+    def __new__(cls, name):
+        instance = super().__new__(cls, 0)
+        instance.name = name
+        return instance
+
+    def __repr__(self):
+        return f"int_sentinel.{self.name}"
+
+    def __hash__(self):
+        return hash((int(self), self.name))
+
+    def __eq__(self, other):
+        return self is other
+
+    __str__ = __reduce__ = __repr__
+
+
+class _IntSentinel:
+    def __init__(self):
+        self._sentinels = {}
+
+    def __getattr__(self, name):
+        if name == "__bases__":
+            raise AttributeError
+        return self._sentinels.setdefault(name, _IntSentinelObject(name))
+
+    def __reduce__(self):
+        return "int_sentinel"
+
+
+int_sentinel = _IntSentinel()

tests/test_tuya_mcu.py

@@ -4,9 +4,14 @@
 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
 from zigpy.zcl import foundation
+from zigpy.zcl.clusters.general import Basic
 
-from tests.common import ClusterListener, MockDatetime
+from tests.common import ClusterListener, MockDatetime, wait_for_zigpy_tasks
 import zhaquirks
 from zhaquirks.tuya import TUYA_MCU_VERSION_RSP, TUYA_SET_TIME, TuyaDPType
 from zhaquirks.tuya.mcu import (
@@ -15,8 +20,15 @@
     TuyaAttributesCluster,
     TuyaClusterData,
     TuyaMCUCluster,
+    TuyaOnOffNM,
+    TuyaQuirkBuilder,
+    TuyaSoilMoisture,
+    TuyaTemperatureMeasurement,
+    TuyaValveWaterConsumed,
 )
 
+from .async_mock import sentinel
+
 zhaquirks.setup()
 
 ZCL_TUYA_VERSION_RSP = b"\x09\x06\x11\x01\x6d\x82"
@@ -375,3 +387,130 @@ async def test_tuya_mcu_classes():
         TuyaClusterData(manufacturer="xiaomi")
     with pytest.raises(ValueError):
         TuyaClusterData(manufacturer=b"")
+
+
+@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
+
+
+async def test_tuya_quirkbuilder(device_mock):
+    """Test adding a v2 Tuya Quirk to the registry and getting back a quirked device."""
+
+    registry = DeviceRegistry()
+
+    class TestEnum(t.enum8):
+        """Test Enum."""
+
+        A = 0x00
+        B = 0x01
+
+    entry = (
+        TuyaQuirkBuilder(device_mock.manufacturer, device_mock.model, registry=registry)
+        .tuya_battery(dp_id=1)
+        .tuya_metering(dp_id=2)
+        .tuya_onoff(dp_id=3)
+        .tuya_soil_moisture(dp_id=4)
+        .tuya_temperature(dp_id=5)
+        .tuya_switch(
+            dp_id=6,
+            attribute_name="test_onoff",
+        )
+        .tuya_number(
+            dp_id=7,
+            attribute_name="test_number",
+            type=t.uint16_t,
+        )
+        .tuya_binary_sensor(
+            dp_id=8,
+            attribute_name="test_binary",
+        )
+        .tuya_sensor(
+            dp_id=9,
+            attribute_name="test_sensor",
+            type=t.uint8_t,
+        )
+        .tuya_enum(
+            dp_id=10,
+            attribute_name="test_enum",
+            enum_class=TestEnum,
+        )
+        .add_to_registry()
+    )
+
+    # coverage for overridden __eq__ method
+    assert entry.adds_metadata[0] != entry.adds_metadata[1]
+    assert entry.adds_metadata[0] != entry
+
+    quirked = registry.get_device(device_mock)
+    assert isinstance(quirked, CustomDeviceV2)
+    assert quirked in registry
+
+    ep = quirked.endpoints[1]
+
+    assert ep.basic is not None
+    assert isinstance(ep.basic, Basic)
+
+    assert ep.tuya_manufacturer is not None
+    assert isinstance(ep.tuya_manufacturer, TuyaMCUCluster)
+
+    tuya_cluster = ep.tuya_manufacturer
+    tuya_listener = ClusterListener(tuya_cluster)
+    assert tuya_cluster.attributes_by_name["test_onoff"].id == 0xEF06
+    assert tuya_cluster.attributes_by_name["test_number"].id == 0xEF07
+    assert tuya_cluster.attributes_by_name["test_binary"].id == 0xEF08
+    assert tuya_cluster.attributes_by_name["test_sensor"].id == 0xEF09
+    assert tuya_cluster.attributes_by_name["test_enum"].id == 0xEF0A
+
+    assert ep.smartenergy_metering is not None
+    assert isinstance(ep.smartenergy_metering, TuyaValveWaterConsumed)
+
+    assert ep.on_off is not None
+    assert isinstance(ep.on_off, TuyaOnOffNM)
+
+    assert ep.soil_moisture is not None
+    assert isinstance(ep.soil_moisture, TuyaSoilMoisture)
+
+    assert ep.temperature is not None
+    assert isinstance(ep.temperature, TuyaTemperatureMeasurement)
+
+    with mock.patch.object(
+        tuya_cluster.endpoint, "request", return_value=foundation.Status.SUCCESS
+    ) as m1:
+        (status,) = await tuya_cluster.write_attributes(
+            {
+                "test_enum": 0x01,
+            }
+        )
+
+        await wait_for_zigpy_tasks()
+        m1.assert_called_with(
+            cluster=61184,
+            sequence=1,
+            data=b"\x01\x01\x00\x00\x01\n\x04\x00\x01\x01",
+            command_id=0,
+            timeout=5,
+            expect_reply=False,
+            use_ieee=False,
+            ask_for_ack=None,
+        )
+        assert status == [
+            foundation.WriteAttributesStatusRecord(foundation.Status.SUCCESS)
+        ]
+
+    assert tuya_listener.attribute_updates[0][0] == 0xEF0A
+    assert tuya_listener.attribute_updates[0][1] == TestEnum.B

zhaquirks/tuya/mcu/__init__.py

@@ -780,7 +780,7 @@ def __init__(
         self.new_attributes: dict[int, foundation.ZCLAttributeDef] = {}
         super().__init__(manufacturer, model, registry)
 
-    def battery(
+    def tuya_battery(
         self,
         dp_id: int,
         power_cfg: PowerConfiguration = TuyaPowerConfigurationCluster2AAA,
@@ -796,7 +796,7 @@ def battery(
         self.adds(power_cfg)
         return self
 
-    def metering(
+    def tuya_metering(
         self,
         dp_id: int,
         metering_cfg: TuyaLocalCluster = TuyaValveWaterConsumed,
@@ -808,8 +808,9 @@ def metering(
             "current_summ_delivered",
         )
         self.adds(metering_cfg)
+        return self
 
-    def onoff(
+    def tuya_onoff(
         self,
         dp_id: int,
         onoff_cfg: TuyaLocalCluster = TuyaOnOffNM,
@@ -821,8 +822,9 @@ def onoff(
             "on_off",
         )
         self.adds(onoff_cfg)
+        return self
 
-    def soil_moisture(
+    def tuya_soil_moisture(
         self,
         dp_id: int,
         soil_cfg: TuyaLocalCluster = TuyaSoilMoisture,
@@ -838,7 +840,7 @@ def soil_moisture(
         self.adds(soil_cfg)
         return self
 
-    def temperature(
+    def tuya_temperature(
         self,
         dp_id: int,
         temp_cfg: TuyaLocalCluster = TuyaTemperatureMeasurement,

zhaquirks/tuya/ts0601_sensor.py

@@ -330,9 +330,9 @@ class TuyaTempHumiditySensorVar04(CustomDevice):
 (
     TuyaQuirkBuilder("_TZE284_aao3yzhs", "TS0601")
     .also_applies_to("_TZE284_sgabhwa6", "TS0601")
-    .temperature(dp_id=5, scale=10)
-    .battery(dp_id=15)
-    .soil_moisture(dp_id=3)
+    .tuya_temperature(dp_id=5, scale=10)
+    .tuya_battery(dp_id=15)
+    .tuya_soil_moisture(dp_id=3)
     .add_to_registry()
 )
 
@@ -342,8 +342,8 @@ class TuyaTempHumiditySensorVar04(CustomDevice):
     .also_applies_to("_TZE200_ga1maeof", "TS0601")
     .also_applies_to("_TZE200_9cqcpkgb", "TS0601")
     .also_applies_to("_TZE204_myd45weu", "TS0601")
-    .temperature(dp_id=5)
-    .battery(dp_id=15)
-    .soil_moisture(dp_id=3)
+    .tuya_temperature(dp_id=5)
+    .tuya_battery(dp_id=15)
+    .tuya_soil_moisture(dp_id=3)
     .add_to_registry()
 )