Add custom device class for Tuya TY0201 sensor

tests/test_tuya.py

@@ -2,6 +2,7 @@
 
 import base64
 import datetime
+import logging
 import struct
 from typing import Final
 from unittest import mock
@@ -14,7 +15,7 @@
 from zigpy.quirks import CustomDevice, get_device
 import zigpy.types as t
 from zigpy.zcl import foundation
-from zigpy.zcl.clusters.general import PowerConfiguration
+from zigpy.zcl.clusters.general import Basic, PowerConfiguration
 from zigpy.zcl.clusters.security import IasZone, ZoneStatus
 from zigpy.zcl.foundation import ZCLAttributeDef
 
@@ -2023,3 +2024,31 @@ async def test_ts601_door_sensor(
     attrs = await cluster.read_attributes(attributes=[attribute])
 
     assert attrs[0].get(attribute) == expected_value
+
+
+async def test_ty0201_bad_direction(zigpy_device_from_v2_quirk, caplog):
+    """Test TY0201 quirk dealing with bad ZCL command direction."""
+
+    device = zigpy_device_from_v2_quirk("_TZ3000_zl1kmjqx", "TY0201")
+    listener = ClusterListener(device.endpoints[1].in_clusters[Basic.cluster_id])
+
+    # The device has a bad ZCL header and reports the incorrect direction for commands
+    with caplog.at_level(logging.WARNING):
+        device.packet_received(
+            t.ZigbeePacket(
+                profile_id=260,
+                cluster_id=0,  # Basic cluster
+                src_ep=1,
+                dst_ep=1,
+                data=t.SerializableBytes(bytes.fromhex("00930A00001001")),
+            )
+        )
+
+    # No warning gets logged
+    warning_messages = [
+        record.message for record in caplog.records if record.levelname == "WARNING"
+    ]
+    assert not warning_messages
+
+    # Our matching logic should be forgiving
+    assert listener.attribute_updates == [(0, t.Bool.true)]

zhaquirks/tuya/ty0201.py

@@ -1,14 +1,50 @@
 """Tuya TY0201 temperature and humidity sensor."""
 
-from zhaquirks.tuya import TuyaPowerConfigurationCluster2AA
+import logging
+
+from zigpy.quirks.v2 import CustomDeviceV2
+import zigpy.types as t
+from zigpy.zcl import Cluster, foundation
+
+from zhaquirks.tuya import BaseEnchantedDevice, TuyaPowerConfigurationCluster2AA
 from zhaquirks.tuya.builder import TuyaQuirkBuilder
 
+_LOGGER = logging.getLogger(__name__)
+
+
+class TY0201Device(CustomDeviceV2, BaseEnchantedDevice):
+    """TY0201 device with direction fix and enchantment."""
+
+    def _find_zcl_cluster(
+        self, hdr: foundation.ZCLHeader, packet: t.ZigbeePacket
+    ) -> Cluster:
+        """Find a cluster for the packet."""
+
+        # TY0201 devices seem to be very lax with their ZCL header's `direction` field,
+        # we should try "flipping" it if matching doesn't work normally.
+        try:
+            return super()._find_zcl_cluster_strict(hdr, packet)
+        except KeyError:
+            _LOGGER.debug(
+                "Packet is coming in the wrong direction, swapping direction and trying again",
+            )
+
+            return super()._find_zcl_cluster_strict(
+                hdr.replace(
+                    frame_control=hdr.frame_control.replace(
+                        direction=hdr.frame_control.direction.flip()
+                    )
+                ),
+                packet,
+            )
+
+
 (
     TuyaQuirkBuilder("_TZ3000_bjawzodf", "TY0201")
     .applies_to("_TZ3000_zl1kmjqx", "TY0201")
     .applies_to("_TZ3000_zl1kmjqx", "")
     .replaces(TuyaPowerConfigurationCluster2AA)
-    .tuya_enchantment()
+    .device_class(TY0201Device)
     .skip_configuration()
     .add_to_registry()
 )