update tests

Author: prairiesnpr

tests/test_tuya_sensor.py

@@ -1,13 +1,25 @@
 """Tests for Tuya Sensor quirks."""
 
+from unittest import mock
+
 import pytest
 from zigpy.zcl import foundation
 from zigpy.zcl.clusters.general import Basic, PowerConfiguration
 from zigpy.zcl.clusters.measurement import RelativeHumidity, TemperatureMeasurement
 
+from tests.common import wait_for_zigpy_tasks
 import zhaquirks
-from zhaquirks.tuya import TuyaLocalCluster
+from zhaquirks.tuya import TUYA_MCU_VERSION_RSP, TuyaLocalCluster
 from zhaquirks.tuya.mcu import TuyaMCUCluster
+from zhaquirks.tuya.tuya_sensor import RespondingTuyaMCUCluster
+
+# Temp DP 1, Humidity DP 2, Battery DP 3
+TUYA_TEMP01_HUM02_BAT03 = b"\x09\xe0\x02\x0b\x33\x01\x02\x00\x04\x00\x00\x00\xfd\x02\x02\x00\x04\x00\x00\x00\x47\x03\x02\x00\x04\x00\x00\x00\x01"
+# Temp DP 1, Humidity DP 2, Battery DP 4
+TUYA_TEMP01_HUM02_BAT04 = b"\x09\xe0\x02\x0b\x33\x01\x02\x00\x04\x00\x00\x00\xfd\x02\x02\x00\x04\x00\x00\x00\x47\x04\x02\x00\x04\x00\x00\x00\x01"
+TUYA_USP = b"\x09\xe0\x02\x0b\x33\x01\x02\x00\x04\x00\x00\x00\xfd\x02\x02\x00\x04\x00\x00\x00\x47\xff\x02\x00\x04\x00\x00\x00\x64"
+
+ZCL_TUYA_VERSION_RSP = b"\x09\x06\x11\x01\x6d\x82"
 
 zhaquirks.setup()
 
@@ -75,8 +87,7 @@ async def test_handle_get_data(
         == data.data.datapoints[2].data.payload * 2
     )
 
-    message = b"\x09\xe0\x02\x0b\x33\x01\x02\x00\x04\x00\x00\x00\xfd\x02\x02\x00\x04\x00\x00\x00\x47\xff\x02\x00\x04\x00\x00\x00\x64"
-    hdr, data = ep.tuya_manufacturer.deserialize(message)
+    hdr, data = ep.tuya_manufacturer.deserialize(TUYA_USP)
 
     status = ep.tuya_manufacturer.handle_get_data(data.data)
     assert status == foundation.Status.UNSUPPORTED_ATTRIBUTE
@@ -92,20 +103,20 @@ async def test_handle_get_data(
 
 
 @pytest.mark.parametrize(
-    "model,manuf,rh_scale,temp_scale",
+    "model,manuf,rh_scale,temp_scale,state_rpt",
     [
-        ("_TZE200_yjjdcqsq", "TS0601", 100, 10),
-        ("_TZE200_9yapgbuv", "TS0601", 100, 10),
-        ("_TZE204_yjjdcqsq", "TS0601", 100, 10),
-        ("_TZE200_utkemkbs", "TS0601", 100, 10),
-        ("_TZE204_utkemkbs", "TS0601", 100, 10),
-        ("_TZE204_yjjdcqsq", "TS0601", 100, 10),
-        ("_TZE204_ksz749x8", "TS0601", 100, 10),
-        ("_TZE204_upagmta9", "TS0601", 100, 10),
+        ("_TZE200_yjjdcqsq", "TS0601", 100, 10, TUYA_TEMP01_HUM02_BAT04),
+        ("_TZE200_9yapgbuv", "TS0601", 100, 10, TUYA_TEMP01_HUM02_BAT04),
+        ("_TZE204_yjjdcqsq", "TS0601", 100, 10, TUYA_TEMP01_HUM02_BAT04),
+        ("_TZE200_utkemkbs", "TS0601", 100, 10, TUYA_TEMP01_HUM02_BAT04),
+        ("_TZE204_utkemkbs", "TS0601", 100, 10, TUYA_TEMP01_HUM02_BAT04),
+        ("_TZE204_yjjdcqsq", "TS0601", 100, 10, TUYA_TEMP01_HUM02_BAT04),
+        ("_TZE204_ksz749x8", "TS0601", 100, 10, TUYA_TEMP01_HUM02_BAT04),
+        ("_TZE204_upagmta9", "TS0601", 100, 10, TUYA_TEMP01_HUM02_BAT03),
     ],
 )
 async def test_handle_get_data_enum_batt(
-    zigpy_device_from_v2_quirk, model, manuf, rh_scale, temp_scale
+    zigpy_device_from_v2_quirk, model, manuf, rh_scale, temp_scale, state_rpt
 ):
     """Test handle_get_data for multiple attributes - enum battery."""
 
@@ -118,11 +129,7 @@ async def test_handle_get_data_enum_batt(
     assert ep.tuya_manufacturer is not None
     assert isinstance(ep.tuya_manufacturer, TuyaMCUCluster)
 
-    if model == "_TZE204_upagmta9":  # Uses dp 3 for battery
-        message = b"\x09\xe0\x02\x0b\x33\x01\x02\x00\x04\x00\x00\x00\xfd\x02\x02\x00\x04\x00\x00\x00\x47\x03\x02\x00\x04\x00\x00\x00\x01"
-    else:
-        message = b"\x09\xe0\x02\x0b\x33\x01\x02\x00\x04\x00\x00\x00\xfd\x02\x02\x00\x04\x00\x00\x00\x47\x04\x02\x00\x04\x00\x00\x00\x01"
-    hdr, data = ep.tuya_manufacturer.deserialize(message)
+    hdr, data = ep.tuya_manufacturer.deserialize(state_rpt)
 
     status = ep.tuya_manufacturer.handle_get_data(data.data)
 
@@ -140,8 +147,7 @@ async def test_handle_get_data_enum_batt(
 
     assert ep.power.get("battery_percentage_remaining") == 100
 
-    message = b"\x09\xe0\x02\x0b\x33\x01\x02\x00\x04\x00\x00\x00\xfd\x02\x02\x00\x04\x00\x00\x00\x47\xff\x02\x00\x04\x00\x00\x00\x64"
-    hdr, data = ep.tuya_manufacturer.deserialize(message)
+    hdr, data = ep.tuya_manufacturer.deserialize(TUYA_USP)
 
     status = ep.tuya_manufacturer.handle_get_data(data.data)
     assert status == foundation.Status.UNSUPPORTED_ATTRIBUTE
@@ -168,3 +174,32 @@ def test_valid_attributes(zigpy_device_from_v2_quirk):
     assert {temperature_attr_id} == temperature_cluster._VALID_ATTRIBUTES
     assert {humidity_attr_id} == humidity_cluster._VALID_ATTRIBUTES
     assert {power_attr_id} == power_config_cluster._VALID_ATTRIBUTES
+
+
+async def test_tuya_version(zigpy_device_from_v2_quirk):
+    """Test TUYA_MCU_VERSION_RSP messages, ensure response."""
+
+    quirked = zigpy_device_from_v2_quirk("_TZE204_upagmta9", "TS0601")
+    ep = quirked.endpoints[1]
+
+    tuya_cluster = ep.tuya_manufacturer
+
+    assert ep.tuya_manufacturer is not None
+    assert isinstance(ep.tuya_manufacturer, TuyaMCUCluster)
+    assert isinstance(ep.tuya_manufacturer, RespondingTuyaMCUCluster)
+
+    # simulate a TUYA_MCU_VERSION_RSP message
+    hdr, args = tuya_cluster.deserialize(ZCL_TUYA_VERSION_RSP)
+    assert hdr.command_id == TUYA_MCU_VERSION_RSP
+
+    with mock.patch.object(
+        ep.tuya_manufacturer._endpoint,
+        "request",
+        return_value=foundation.Status.SUCCESS,
+    ) as m1:
+        ep.tuya_manufacturer.handle_message(hdr, args)
+        await wait_for_zigpy_tasks()
+
+        res_hdr = foundation.ZCLHeader.deserialize(m1.await_args[1]["data"])
+        assert not res_hdr[0].manufacturer
+        assert not res_hdr[0].frame_control.is_manufacturer_specific