Add Enchanted Device method to TuyaQuirkBuilder (#3661)

tests/test_tuya_builder.py

@@ -12,6 +12,7 @@
 
 from tests.common import ClusterListener, wait_for_zigpy_tasks
 import zhaquirks
+from zhaquirks.tuya import TUYA_QUERY_DATA
 from zhaquirks.tuya.builder import (
     TuyaIasContact,
     TuyaIasFire,
@@ -187,3 +188,73 @@ class TestEnum(t.enum8):
 
     assert tuya_listener.attribute_updates[0][0] == 0xEF0A
     assert tuya_listener.attribute_updates[0][1] == TestEnum.B
+
+
+@pytest.mark.parametrize(
+    "read_attr_spell,data_query_spell",
+    [
+        (True, False),
+        (False, True),
+        (True, True),
+        (False, False),
+    ],
+)
+async def test_tuya_spell(device_mock, read_attr_spell, data_query_spell):
+    """Test that enchanted Tuya devices have their spells applied during configuration."""
+    registry = DeviceRegistry()
+
+    entry = (
+        TuyaQuirkBuilder(device_mock.manufacturer, device_mock.model, registry=registry)
+        .tuya_battery(dp_id=1)
+        .tuya_onoff(dp_id=3)
+        .tuya_enchantment(
+            read_attr_spell=read_attr_spell, data_query_spell=data_query_spell
+        )
+        .skip_configuration()
+        .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
+
+    request_patch = mock.patch("zigpy.zcl.Cluster.request", mock.AsyncMock())
+    with request_patch as request_mock:
+        request_mock.return_value = (foundation.Status.SUCCESS, "done")
+
+        # call apply_custom_configuration() on each EnchantedDevice
+        # ZHA does this during device configuration normally
+        await quirked.apply_custom_configuration()
+
+        # the number of Tuya spells that are allowed to be cast, so the sum of enabled Tuya spells
+        enabled_tuya_spells_num = (
+            quirked.tuya_spell_read_attributes + quirked.tuya_spell_data_query
+        )
+
+        # verify request was called the correct number of times
+        assert request_mock.call_count == enabled_tuya_spells_num
+
+        # used to check list of mock calls below
+        messages = 0
+
+        # check 'attribute read spell' was cast correctly (if enabled)
+        if quirked.tuya_spell_read_attributes:
+            assert (
+                request_mock.mock_calls[messages][1][1]
+                == foundation.GeneralCommand.Read_Attributes
+            )
+            assert request_mock.mock_calls[messages][1][3] == [4, 0, 1, 5, 7, 65534]
+            messages += 1
+
+        # check 'query data spell' was cast correctly (if enabled)
+        if quirked.tuya_spell_data_query:
+            assert not request_mock.mock_calls[messages][1][0]
+            assert request_mock.mock_calls[messages][1][1] == TUYA_QUERY_DATA
+            messages += 1
+
+        request_mock.reset_mock()

zhaquirks/tuya/__init__.py

@@ -7,7 +7,7 @@
 import logging
 from typing import Any, Optional, Union
 
-from zigpy.quirks import CustomCluster, CustomDevice
+from zigpy.quirks import BaseCustomDevice, CustomCluster, CustomDevice
 import zigpy.types as t
 from zigpy.zcl import BaseAttributeDefs, foundation
 from zigpy.zcl.clusters.closures import WindowCovering
@@ -529,7 +529,7 @@ async def write_attributes(self, attributes, manufacturer=None):
         return [[foundation.WriteAttributesStatusRecord(foundation.Status.SUCCESS)]]
 
 
-class EnchantedDevice(CustomDevice):
+class BaseEnchantedDevice(BaseCustomDevice):
     """Class for Tuya devices which need to be unlocked by casting a 'spell'.
 
     The spell is applied during device configuration.
@@ -570,6 +570,10 @@ async def spell_data_query(self):
         self.debug("Executed data query spell on Tuya device %s", self.ieee)
 
 
+class EnchantedDevice(CustomDevice, BaseEnchantedDevice):
+    """Enchanted device class for v1 quirks."""
+
+
 class TuyaOnOff(CustomCluster, OnOff):
     """Tuya On/Off cluster for On/Off device."""
 

zhaquirks/tuya/builder/__init__.py

@@ -6,7 +6,7 @@
 
 from zigpy.quirks import _DEVICE_REGISTRY
 from zigpy.quirks.registry import DeviceRegistry
-from zigpy.quirks.v2 import QuirkBuilder, QuirksV2RegistryEntry
+from zigpy.quirks.v2 import CustomDeviceV2, QuirkBuilder, QuirksV2RegistryEntry
 from zigpy.quirks.v2.homeassistant import EntityPlatform, EntityType
 from zigpy.quirks.v2.homeassistant.binary_sensor import BinarySensorDeviceClass
 from zigpy.quirks.v2.homeassistant.number import NumberDeviceClass
@@ -23,6 +23,7 @@
 
 from zhaquirks.tuya import (
     TUYA_CLUSTER_ID,
+    BaseEnchantedDevice,
     PowerConfiguration,
     TuyaLocalCluster,
     TuyaPowerConfigurationCluster2AAA,
@@ -522,6 +523,21 @@ def tuya_sensor(
 
         return self
 
+    def tuya_enchantment(
+        self, read_attr_spell: bool = True, data_query_spell: bool = False
+    ) -> QuirkBuilder:
+        """Set the Tuya enchantment spells."""
+
+        class EnchantedDeviceV2(CustomDeviceV2, BaseEnchantedDevice):
+            """Enchanted device class for v2 quirks."""
+
+        EnchantedDeviceV2.tuya_spell_read_attributes = read_attr_spell
+        EnchantedDeviceV2.tuya_spell_data_query = data_query_spell
+
+        self.device_class(EnchantedDeviceV2)
+
+        return self
+
     def add_to_registry(self) -> QuirksV2RegistryEntry:
         """Build the quirks v2 registry entry."""