Refactor Tuya spell implementation, add Tuya query data spell (#2940)

Author: TheJulianJES

tests/test_tuya.py

@@ -17,6 +17,7 @@
     GreenPowerProxy,
     Groups,
     Identify,
+    OnOff,
     Ota,
     PowerConfiguration,
 )
@@ -36,10 +37,13 @@
     ZONE_STATUS_CHANGE_COMMAND,
 )
 from zhaquirks.tuya import (
+    TUYA_QUERY_DATA,
     Data,
+    EnchantedDevice,
     TuyaEnchantableCluster,
     TuyaManufClusterAttributes,
     TuyaNewManufCluster,
+    TuyaZBOnOffAttributeCluster,
 )
 import zhaquirks.tuya.sm0202_motion
 import zhaquirks.tuya.ts011f_plug
@@ -1608,22 +1612,58 @@ async def test_power_config_no_bind(zigpy_device_from_quirk, quirk):
         assert len(bind_mock.mock_calls) == 0
 
 
-ENCHANTED_QUIRKS = []
+class TuyaTestSpellDevice(EnchantedDevice):
+    """Tuya test spell device."""
+
+    tuya_spell_data_query = True  # enable additional data query spell
+
+    signature = {
+        MODELS_INFO: [("UjqjHq6ZErY23tgs", "zo9WD7q5dbvDj96y")],
+        ENDPOINTS: {
+            1: {
+                PROFILE_ID: zha.PROFILE_ID,
+                DEVICE_TYPE: zha.DeviceType.SMART_PLUG,
+                INPUT_CLUSTERS: [
+                    Basic.cluster_id,
+                    OnOff.cluster_id,
+                    TuyaNewManufCluster.cluster_id,
+                ],
+                OUTPUT_CLUSTERS: [],
+            }
+        },
+    }
+
+    replacement = {
+        ENDPOINTS: {
+            1: {
+                PROFILE_ID: zha.PROFILE_ID,
+                DEVICE_TYPE: zha.DeviceType.SMART_PLUG,
+                INPUT_CLUSTERS: [
+                    Basic.cluster_id,
+                    TuyaZBOnOffAttributeCluster,
+                    TuyaNewManufCluster,
+                ],
+                OUTPUT_CLUSTERS: [],
+            }
+        }
+    }
+
+
+ENCHANTED_QUIRKS = [TuyaTestSpellDevice]
 for manufacturer in zigpy.quirks._DEVICE_REGISTRY._registry.values():
     for model_quirk_list in manufacturer.values():
         for quirk_entry in model_quirk_list:
             if quirk_entry in ENCHANTED_QUIRKS:
                 continue
-            # right now, this basically includes `issubclass(quirk, EnchantedDevice)`, as that sets `TUYA_SPELL`
-            if getattr(quirk_entry, "TUYA_SPELL", False):
+            if issubclass(quirk_entry, EnchantedDevice):
                 ENCHANTED_QUIRKS.append(quirk_entry)
 
 del quirk_entry, model_quirk_list, manufacturer
 
 
 @mock.patch("zigpy.zcl.Cluster.bind", mock.AsyncMock())
 async def test_tuya_spell(zigpy_device_from_quirk):
-    """Test that enchanted Tuya devices have their spell applied when binding OnOff cluster."""
+    """Test that enchanted Tuya devices have their spell applied when binding bindable cluster."""
     non_bindable_cluster_ids = [
         Basic.cluster_id,
         Identify.cluster_id,
@@ -1639,6 +1679,7 @@ async def test_tuya_spell(zigpy_device_from_quirk):
 
         for quirk in ENCHANTED_QUIRKS:
             device = zigpy_device_from_quirk(quirk)
+            assert isinstance(device, EnchantedDevice)
 
             # fail if SKIP_CONFIGURATION is set, as that will cause ZHA to not call bind()
             if getattr(device, SKIP_CONFIGURATION, False):
@@ -1661,45 +1702,84 @@ async def test_tuya_spell(zigpy_device_from_quirk):
                 ):
                     await cluster.bind()
 
-            # check that exactly one Tuya spell was cast
+            # the number of Tuya spells that are allowed to be cast, so the sum of enabled Tuya spells
+            enabled_tuya_spells_num = (
+                device.tuya_spell_read_attributes + device.tuya_spell_data_query
+            )
+
+            # skip if no Tuya spells are enabled,
+            # this case is already handled in the test_tuya_spell_devices_valid() test
+            if enabled_tuya_spells_num == 0:
+                continue
+
+            # check that exactly a Tuya spell was cast
             if len(request_mock.mock_calls) == 0:
                 pytest.fail(
                     f"Enchanted quirk {quirk} did not cast a Tuya spell. "
                     f"One bindable cluster subclassing `TuyaEnchantableCluster` on endpoint 1 needs to be implemented. "
                     f"Also check that enchanted bindable clusters do not modify their `ep_attribute`, "
                     f"as ZHA will not call bind() in that case."
                 )
-            elif len(request_mock.mock_calls) > 1:
+            # check that no more than one call was made for each enabled spell
+            elif len(request_mock.mock_calls) > enabled_tuya_spells_num:
                 pytest.fail(
                     f"Enchanted quirk {quirk} cast more than one Tuya spell. "
                     f"Make sure to only implement one cluster subclassing `TuyaEnchantableCluster` on endpoint 1."
                 )
 
-            assert (
-                request_mock.mock_calls[0][1][1]
-                == foundation.GeneralCommand.Read_Attributes
-            )  # read attributes
-            assert request_mock.mock_calls[0][1][3] == [4, 0, 1, 5, 7, 65534]  # spell
+            # used to check list of mock calls below
+            messages = 0
+
+            # check 'attribute read spell' was cast correctly (if enabled)
+            if device.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 device.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()
 
 
 def test_tuya_spell_devices_valid():
     """Test that all enchanted Tuya devices implement at least one enchanted cluster."""
 
     for quirk in ENCHANTED_QUIRKS:
-        enchanted_clusters = 0
+        # check that at least one Tuya spell is enabled for an EnchantedDevice
+        if not quirk.tuya_spell_read_attributes and not quirk.tuya_spell_data_query:
+            pytest.fail(
+                f"Enchanted quirk {quirk} does not have any Tuya spells enabled. "
+                f"Enable at least one Tuya spell by setting `TUYA_SPELL_READ_ATTRIBUTES` or `TUYA_SPELL_DATA_QUERY` "
+                f"or inherit CustomDevice rather than EnchantedDevice."
+            )
+
+        enchanted_clusters = 0  # number of clusters subclassing TuyaEnchantableCluster
+        tuya_cluster_exists = False  # cluster subclassing TuyaNewManufCluster existing
 
         # iterate over all clusters in the replacement
         for endpoint_id, endpoint in quirk.replacement[ENDPOINTS].items():
             if endpoint_id != 1:  # spell is only activated on endpoint 1 for now
                 continue
             for cluster in endpoint[INPUT_CLUSTERS] + endpoint[OUTPUT_CLUSTERS]:
-                if not isinstance(cluster, int) and issubclass(
-                    cluster, TuyaEnchantableCluster
-                ):
-                    enchanted_clusters += 1
-
-        # one EnchantedDevice must have exactly one enchanted cluster on endpoint 1
+                if not isinstance(cluster, int):
+                    # count all clusters which would apply the spell on bind()
+                    if issubclass(cluster, TuyaEnchantableCluster):
+                        enchanted_clusters += 1
+                    # check if there's a valid Tuya cluster where the id wasn't modified
+                    if (
+                        issubclass(cluster, TuyaNewManufCluster)
+                        and cluster.cluster_id == TuyaNewManufCluster.cluster_id
+                    ):
+                        tuya_cluster_exists = True
+
+        # an EnchantedDevice must have exactly one enchanted cluster on endpoint 1
         if enchanted_clusters == 0:
             pytest.fail(
                 f"{quirk} does not have a cluster subclassing `TuyaEnchantableCluster` on endpoint 1 "
@@ -1709,3 +1789,9 @@ def test_tuya_spell_devices_valid():
             pytest.fail(
                 f"{quirk} has more than one cluster subclassing `TuyaEnchantableCluster` on endpoint 1"
             )
+
+        # an EnchantedDevice with the data query spell must also have a cluster subclassing TuyaNewManufCluster
+        if quirk.tuya_spell_data_query and not tuya_cluster_exists:
+            pytest.fail(
+                f"{quirk} set Tuya data query spell but has no cluster subclassing `TuyaNewManufCluster` on endpoint 1"
+            )

zhaquirks/tuya/__init__.py

@@ -9,7 +9,7 @@
 import zigpy.types as t
 from zigpy.zcl import foundation
 from zigpy.zcl.clusters.closures import WindowCovering
-from zigpy.zcl.clusters.general import LevelControl, OnOff, PowerConfiguration
+from zigpy.zcl.clusters.general import Basic, LevelControl, OnOff, PowerConfiguration
 from zigpy.zcl.clusters.homeautomation import ElectricalMeasurement
 from zigpy.zcl.clusters.hvac import Thermostat, UserInterface
 from zigpy.zcl.clusters.smartenergy import Metering
@@ -527,6 +527,18 @@ async def write_attributes(self, attributes, manufacturer=None):
         return [[foundation.WriteAttributesStatusRecord(foundation.Status.SUCCESS)]]
 
 
+class EnchantedDevice(CustomDevice):
+    """Class for Tuya devices which need to be unlocked by casting a 'spell'. This happens during binding.
+
+    To make sure the spell is cast, the device needs to implement a subclass of `TuyaEnchantableCluster`.
+    For more information, see the documentation of `TuyaEnchantableCluster`.
+    """
+
+    # These values can be overridden from a quirk to enable (or disable) additional Tuya spells:
+    tuya_spell_read_attributes: bool = True  # spell reading attributes on Basic cluster
+    tuya_spell_data_query: bool = False  # additional spell needed for some devices
+
+
 class TuyaEnchantableCluster(CustomCluster):
     """Tuya cluster that casts a magic spell if `TUYA_SPELL` is set.
 
@@ -549,22 +561,44 @@ class TuyaEnchantableCluster(CustomCluster):
 
     async def bind(self):
         """Bind cluster and start casting the spell if necessary."""
-        # check if the device needs to have the spell cast
+        device = self.endpoint.device
+
+        # check if the device is an EnchantedDevice
         # and since the cluster can be used on multiple endpoints, check that it's endpoint 1
-        if (
-            getattr(self.endpoint.device, "TUYA_SPELL", False)
-            and self.endpoint.endpoint_id == 1
-        ):
-            await self.spell()
+        if isinstance(device, EnchantedDevice) and self.endpoint.endpoint_id == 1:
+            if device.tuya_spell_read_attributes:
+                await self.spell_attribute_reads()
+            if device.tuya_spell_data_query:
+                await self.spell_data_query()
+
         return await super().bind()
 
-    async def spell(self):
-        """Cast spell, so the Tuya device works correctly."""
-        self.debug("Executing spell on Tuya device %s", self.endpoint.device.ieee)
+    async def spell_attribute_reads(self):
+        """Cast 'attribute read' spell, so the Tuya device works correctly."""
+        self.debug(
+            "Executing attribute read spell on Tuya device %s",
+            self.endpoint.device.ieee,
+        )
         attr_to_read = [4, 0, 1, 5, 7, 0xFFFE]
-        basic_cluster = self.endpoint.device.endpoints[1].in_clusters[0]
+        basic_cluster = self.endpoint.device.endpoints[1].in_clusters[Basic.cluster_id]
         await basic_cluster.read_attributes(attr_to_read)
-        self.debug("Executed spell on Tuya device %s", self.endpoint.device.ieee)
+        self.debug(
+            "Executed attribute read spell on Tuya device %s", self.endpoint.device.ieee
+        )
+
+    async def spell_data_query(self):
+        """Cast 'data query' spell, also required for some Tuya devices to send data."""
+        self.debug(
+            "Executing data query spell on Tuya device %s", self.endpoint.device.ieee
+        )
+        # tests verify that a device with an enabled 'data query spell' has a TuyaNewManufCluster (subclass)
+        tuya_cluster = self.endpoint.device.endpoints[1].in_clusters[
+            TuyaNewManufCluster.cluster_id
+        ]
+        await tuya_cluster.command(TUYA_QUERY_DATA)
+        self.debug(
+            "Executed data query spell on Tuya device %s", self.endpoint.device.ieee
+        )
 
 
 class TuyaOnOff(TuyaEnchantableCluster, OnOff):

zhaquirks/tuya/mcu/__init__.py

@@ -3,7 +3,6 @@
 import datetime
 from typing import Any, Callable, Dict, Optional, Tuple, Union
 
-from zigpy.quirks import CustomDevice
 import zigpy.types as t
 from zigpy.zcl import foundation
 from zigpy.zcl.clusters.general import LevelControl, OnOff
@@ -24,6 +23,9 @@
     TuyaTimePayload,
 )
 
+# add EnchantedDevice import for custom quirks backwards compatibility
+from zhaquirks.tuya import EnchantedDevice  # noqa: F401
+
 # New manufacturer attributes
 ATTR_MCU_VERSION = 0xEF00
 
@@ -720,13 +722,3 @@ class TuyaLevelControlManufCluster(TuyaMCUCluster):
         17: "_dp_2_attr_update",
         18: "_dp_2_attr_update",
     }
-
-
-class EnchantedDevice(CustomDevice):
-    """Class for Tuya devices which need to be unlocked by casting a 'spell'. This happens during binding.
-
-    To make sure the spell is cast, the device needs to implement a subclass of `TuyaEnchantableCluster`.
-    For more information, see the documentation of `TuyaEnchantableCluster`.
-    """
-
-    TUYA_SPELL = True