Fix switch events for Aqara T2 relay (#2939)

Author: trapacska

tests/test_xiaomi.py

@@ -32,6 +32,8 @@
 
 import zhaquirks
 from zhaquirks.const import (
+    BUTTON_1,
+    BUTTON_2,
     DEVICE_TYPE,
     ENDPOINTS,
     INPUT_CLUSTERS,
@@ -1609,31 +1611,6 @@ async def test_xiaomi_e1_roller_commands_1(zigpy_device_from_quirk, command, val
     assert multistate_cluster._write_attributes.call_args[0][0][0].value.value == value
 
 
-@pytest.mark.parametrize(
-    "quirk",
-    (zhaquirks.xiaomi.aqara.switch_acn047.AqaraT2Relay,),
-)
-async def test_aqara_t2_relay(zigpy_device_from_quirk, quirk):
-    """Test Aqara T2 relay."""
-
-    device = zigpy_device_from_quirk(quirk)
-    mi_cluster = device.endpoints[1].multistate_input
-    mi_listener = ClusterListener(mi_cluster)
-
-    mi_cluster.update_attribute(MultistateInput.AttributeDefs.present_value.id, 1)
-    assert len(mi_listener.attribute_updates) == 1
-    assert mi_listener.attribute_updates[0][0] == 0
-    assert mi_listener.attribute_updates[0][1] == "single"
-
-    mi_cluster.update_attribute(MultistateInput.AttributeDefs.state_text.id, "foo")
-    assert len(mi_listener.attribute_updates) == 2
-    assert (
-        mi_listener.attribute_updates[1][0]
-        == MultistateInput.AttributeDefs.state_text.id
-    )
-    assert mi_listener.attribute_updates[1][1] == "foo"
-
-
 @pytest.mark.parametrize(
     "command, value",
     [
@@ -1664,6 +1641,30 @@ async def test_xiaomi_e1_roller_commands_2(zigpy_device_from_quirk, command, val
     )
 
 
+@pytest.mark.parametrize("endpoint", [(1), (2)])
+async def test_aqara_t2_relay(zigpy_device_from_quirk, endpoint):
+    """Test Aqara T2 relay."""
+
+    device = zigpy_device_from_quirk(zhaquirks.xiaomi.aqara.switch_acn047.AqaraT2Relay)
+    mi_cluster = device.endpoints[endpoint].multistate_input
+    mi_listener = ClusterListener(mi_cluster)
+
+    buttons = {1: BUTTON_1, 2: BUTTON_2}
+
+    mi_cluster.update_attribute(MultistateInput.AttributeDefs.present_value.id, 1)
+    assert len(mi_listener.attribute_updates) == 1
+    assert mi_listener.attribute_updates[0][0] == 0
+    assert mi_listener.attribute_updates[0][1] == buttons[endpoint]
+
+    mi_cluster.update_attribute(MultistateInput.AttributeDefs.state_text.id, "foo")
+    assert len(mi_listener.attribute_updates) == 2
+    assert (
+        mi_listener.attribute_updates[1][0]
+        == MultistateInput.AttributeDefs.state_text.id
+    )
+    assert mi_listener.attribute_updates[1][1] == "foo"
+
+
 def test_aqara_acn003_signature_match(assert_signature_matches_quirk):
     signature = {
         "node_descriptor": "NodeDescriptor(logical_type=<LogicalType.Router: 1>, complex_descriptor_available=0, user_descriptor_available=0, reserved=0, aps_flags=0, frequency_band=<FrequencyBand.Freq2400MHz: 8>, mac_capability_flags=<MACCapabilityFlags.FullFunctionDevice|MainsPowered|RxOnWhenIdle|AllocateAddress: 142>, manufacturer_code=4447, maximum_buffer_size=82, maximum_incoming_transfer_size=82, server_mask=11264, maximum_outgoing_transfer_size=82, descriptor_capability_field=<DescriptorCapability.NONE: 0>, *allocate_address=True, *is_alternate_pan_coordinator=False, *is_coordinator=False, *is_end_device=False, *is_full_function_device=True, *is_mains_powered=True, *is_receiver_on_when_idle=True, *is_router=True, *is_security_capable=False)",

zhaquirks/xiaomi/aqara/switch_acn047.py

@@ -19,8 +19,11 @@
 
 from zhaquirks.const import (
     ATTR_ID,
+    BUTTON_1,
+    BUTTON_2,
     COMMAND,
     DEVICE_TYPE,
+    ENDPOINT_ID,
     ENDPOINTS,
     INPUT_CLUSTERS,
     MODELS_INFO,
@@ -40,8 +43,7 @@
     XiaomiCustomDevice,
 )
 
-PRESS_TYPES = {1: "single"}
-SINGLE = "single"
+PRESS_TYPES = {1: BUTTON_1, 2: BUTTON_2}
 STATUS_TYPE_ATTR = 0x0055  # decimal = 85
 
 
@@ -54,14 +56,15 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
     def _update_attribute(self, attrid, value):
-        if attrid == STATUS_TYPE_ATTR:
-            self._current_state = PRESS_TYPES.get(value)
+        if attrid == STATUS_TYPE_ATTR and value == 1:
+            self._current_state = PRESS_TYPES.get(self.endpoint.endpoint_id)
             event_args = {
                 PRESS_TYPE: self._current_state,
                 ATTR_ID: attrid,
                 VALUE: value,
+                ENDPOINT_ID: self.endpoint.endpoint_id,
             }
-            self.listener_event(ZHA_SEND_EVENT, self, self._current_state, event_args)
+            self.listener_event(ZHA_SEND_EVENT, self._current_state, event_args)
             super()._update_attribute(0, self._current_state)
         else:
             super()._update_attribute(attrid, value)
@@ -206,5 +209,12 @@ class AqaraT2Relay(XiaomiCustomDevice):
     }
 
     device_automation_triggers = {
-        (SHORT_PRESS, SHORT_PRESS): {COMMAND: SINGLE},
+        (SHORT_PRESS, PRESS_TYPES[1]): {
+            ENDPOINT_ID: 1,
+            COMMAND: PRESS_TYPES[1],
+        },
+        (SHORT_PRESS, PRESS_TYPES[2]): {
+            ENDPOINT_ID: 2,
+            COMMAND: PRESS_TYPES[2],
+        },
     }