Merge branch 'dev' into add-SZ-T04

Author: prairiesnpr

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]

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

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
@@ -12,6 +11,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,
@@ -24,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",
     [
@@ -191,3 +170,73 @@ class ModTuyaMCUCluster(TuyaMCUCluster):
 
     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()

tests/test_tuya_motion.py

@@ -24,6 +24,7 @@
     "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),

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

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, mcu_cluster=TuyaMCUCluster) -> QuirksV2RegistryEntry:
         """Build the quirks v2 registry entry."""
 

zhaquirks/tuya/ts0601_motion.py

@@ -126,6 +126,7 @@ class TuyaSelfCheckResult(t.enum8):
 (
     base_tuya_motion.clone()
     .applies_to("_TZE200_ya4ft0w4", "TS0601")
+    .applies_to("_TZE204_ya4ft0w4", "TS0601")
     .tuya_dp(
         dp_id=1,
         ep_attribute=TuyaOccupancySensing.ep_attribute,