Repository: zigpy/zha-device-handlers
Pull request #3661 (2024-12-31 19:46:09)

Add Enchanted Device method to TuyaQuirkBuilder
## Proposed change
<!--
  Explain your proposed change below.
-->

This enables adding enchanted devices using TuyaQuirkBuilder.

```python
(
        TuyaQuirkBuilder("manufacturer", "model")
        .tuya_battery(dp_id=1)
        .tuya_onoff(dp_id=3)
        .tuya_enchantment(
            read_attr_spell=True, data_query_spell=False
        )
        .skip_configuration()
        .add_to_registry()
)
```

## Additional information
<!--
  Please include any additional information that is important to this PR.
  For example, if this PR is a potentially breaking change, mention that here.
  If this PR requires other PRs to be merged in HA Core or other projects, mention that.
  Lastly, if this PR fixes a specific issue, please include "Fixes #xxxx".
-->


## Checklist
<!--
  Put an 'x' in all boxes that apply.
  Note: You do not need to tick all boxes before creating a PR.
-->

- [x] The changes are tested and work correctly
- [x] `pre-commit` checks pass / the code has been formatted using Black
- [x] Tests have been added to verify that the new code works


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/zigpy/zha-device-handlers/pull/3661?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zigpy) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 89.87%. Comparing base [(`8124411`)](https://app.codecov.io/gh/zigpy/zha-device-handlers/commit/812441104427e74e13a8168ea916597f87534d72?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zigpy) to head [(`d954b06`)](https://app.codecov.io/gh/zigpy/zha-device-handlers/commit/d954b0642d860035b461049008bf1b527c7b7bab?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zigpy).

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##              dev    #3661      +/-   ##
==========================================
+ Coverage   89.85%   89.87%   +0.02%     
==========================================
  Files         322      322              
  Lines       10379    10405      +26     
==========================================
+ Hits         9326     9352      +26     
  Misses       1053     1053              
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/zigpy/zha-device-handlers/pull/3661?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zigpy).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zigpy).


Patch:
diff --git a/pyproject.toml b/pyproject.toml
index be77dd2..bd97df9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -14,7 +14,7 @@ readme = "README.md"
 license = {text = "Apache-2.0"}
 requires-python = ">=3.12"
 dependencies = [
-    "zigpy>=0.70.0",
+    "zigpy>=0.74.0",
 ]
 
 [tool.setuptools.packages.find]
diff --git a/zhaquirks/tuya/__init__.py b/zhaquirks/tuya/__init__.py
index ac4554f..b63b48b 100644
--- a/zhaquirks/tuya/__init__.py
+++ b/zhaquirks/tuya/__init__.py
@@ -7,7 +7,7 @@ import enum
 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 @@ class TuyaManufClusterAttributes(TuyaManufCluster):
         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 @@ class EnchantedDevice(CustomDevice):
         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."""
 
diff --git a/zhaquirks/tuya/builder/__init__.py b/zhaquirks/tuya/builder/__init__.py
index 1c45989..b683add 100644
--- a/zhaquirks/tuya/builder/__init__.py
+++ b/zhaquirks/tuya/builder/__init__.py
@@ -6,7 +6,7 @@ from typing import Any, Optional
 
 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 zigpy.zcl.clusters.smartenergy import Metering
 
 from zhaquirks.tuya import (
     TUYA_CLUSTER_ID,
+    BaseEnchantedDevice,
     PowerConfiguration,
     TuyaLocalCluster,
     TuyaPowerConfigurationCluster2AAA,
@@ -522,6 +523,21 @@ class TuyaQuirkBuilder(QuirkBuilder):
 
         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."""
 
diff --git a/zhaquirks/tuya/ts0601_motion.py b/zhaquirks/tuya/ts0601_motion.py
index 2c12194..65ac852 100644
--- a/zhaquirks/tuya/ts0601_motion.py
+++ b/zhaquirks/tuya/ts0601_motion.py
@@ -126,6 +126,7 @@ base_tuya_motion = (
 (
     base_tuya_motion.clone()
     .applies_to("_TZE200_ya4ft0w4", "TS0601")
+    .applies_to("_TZE204_ya4ft0w4", "TS0601")
     .tuya_dp(
         dp_id=1,
         ep_attribute=TuyaOccupancySensing.ep_attribute,
diff --git a/zhaquirks/xiaomi/aqara/remote_h1.py b/zhaquirks/xiaomi/aqara/remote_h1.py
index 786bd15..e15ab52 100644
--- a/zhaquirks/xiaomi/aqara/remote_h1.py
+++ b/zhaquirks/xiaomi/aqara/remote_h1.py
@@ -1,6 +1,7 @@
 """Aqara H1-series wireless remote."""
 
 from zigpy import types
+from zigpy.profiles import zha
 from zigpy.quirks.v2 import ClusterType, QuirkBuilder
 from zigpy.zcl.clusters.general import Identify, OnOff
 from zigpy.zcl.foundation import BaseAttributeDefs, DataTypeId, ZCLAttributeDef
@@ -144,6 +145,9 @@ class PowerConfigurationClusterH1Remote(PowerConfigurationCluster):
     # .friendly_name(
     #     manufacturer="Aqara", model="Wireless Remote Switch H1 (Double Rocker)"
     # )
+    # add endpoints that are present but not exposed sometimes
+    .adds_endpoint(2, device_type=zha.DeviceType.ON_OFF_LIGHT_SWITCH)
+    .adds_endpoint(3, device_type=zha.DeviceType.ON_OFF_LIGHT_SWITCH)
     .replaces(AqaraRemoteManuSpecificCluster)
     .adds(Identify)
     .replaces(MultistateInputCluster)


Test patch:
diff --git a/requirements_test.txt b/requirements_test.txt
index 0159e5f..417a0ce 100644
--- a/requirements_test.txt
+++ b/requirements_test.txt
@@ -13,5 +13,5 @@ pytest-sugar
 pytest-timeout
 pytest-asyncio
 pytest>=7.1.3
-zigpy>=0.70
+zigpy>=0.74.0
 ruff==0.0.261
diff --git a/tests/test_tuya_builder.py b/tests/test_tuya_builder.py
index a188348..2f25115 100644
--- a/tests/test_tuya_builder.py
+++ b/tests/test_tuya_builder.py
@@ -12,6 +12,7 @@ from zigpy.zcl.clusters.general import Basic
 
 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 @@ async def test_tuya_quirkbuilder(device_mock):
 
     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()
diff --git a/tests/test_tuya_motion.py b/tests/test_tuya_motion.py
index 10e1d08..a11f9d2 100644
--- a/tests/test_tuya_motion.py
+++ b/tests/test_tuya_motion.py
@@ -24,6 +24,7 @@ zhaquirks.setup()
     "model,manuf,occ_msg",
     [
         ("_TZE200_ya4ft0w4", "TS0601", ZCL_TUYA_MOTION),
+        ("_TZE204_ya4ft0w4", "TS0601", ZCL_TUYA_MOTION),
         ("_TZE200_7hfcudw5", "TS0601", ZCL_TUYA_MOTION_V2),
         ("_TZE200_ppuj1vem", "TS0601", ZCL_TUYA_MOTION_V2),
         ("_TZE200_ar0slwnd", "TS0601", ZCL_TUYA_MOTION),
diff --git a/tests/test_xiaomi.py b/tests/test_xiaomi.py
index ef2c7ad..811a667 100644
--- a/tests/test_xiaomi.py
+++ b/tests/test_xiaomi.py
@@ -1804,3 +1804,20 @@ async def test_aqara_fp1e_sensor(
     assert len(ias_listener.attribute_updates) == 1
     assert ias_listener.attribute_updates[0][0] == IasZone.AttributeDefs.zone_status.id
     assert ias_listener.attribute_updates[0][1] == expected_motion_status
+
+
+def test_h1_wireless_remotes(zigpy_device_from_v2_quirk):
+    """Test Aqara H1 wireless remote quirk adds missing endpoints."""
+    # create device with endpoint 1 only and verify we don't get a KeyError
+    quirk = zigpy_device_from_v2_quirk(LUMI, "lumi.remote.b28ac1")
+
+    # verify the quirk adds endpoints 2 and 3
+    assert 2 in quirk.endpoints
+    assert 3 in quirk.endpoints
+
+    # verify the quirk adds the correct clusters to the new endpoints
+    assert OnOff.cluster_id in quirk.endpoints[2].out_clusters
+    assert OnOff.cluster_id in quirk.endpoints[3].out_clusters
+
+    assert MultistateInput.cluster_id in quirk.endpoints[2].in_clusters
+    assert MultistateInput.cluster_id in quirk.endpoints[3].in_clusters