Improve Legrand 064882 cable outlet heat mode

tests/test_legrand.py

@@ -1,8 +1,13 @@
 """Tests for Legrand."""
 
+from unittest import mock
+
 import pytest
+from zigpy.zcl import foundation
 
+from tests.common import ClusterListener
 import zhaquirks
+from zhaquirks.legrand import LEGRAND
 
 zhaquirks.setup()
 
@@ -56,3 +61,79 @@ def test_light_switch_with_neutral_signature(assert_signature_matches_quirk):
     assert_signature_matches_quirk(
         zhaquirks.legrand.switch.LightSwitchWithNeutral, signature
     )
+
+
+async def test_legrand_wire_pilot_cluster_write_attrs(zigpy_device_from_v2_quirk):
+    """Test Legrand cable outlet heat mode attr writing."""
+
+    device = zigpy_device_from_v2_quirk(f" {LEGRAND}", " Cable outlet")
+    legrand_wire_pilot_cluster = device.endpoints[1].legrand_wire_pilot_cluster
+    legrand_wire_pilot_cluster._write_attributes = mock.AsyncMock()
+    legrand_wire_pilot_cluster.set_heat_mode = mock.AsyncMock()
+
+    await legrand_wire_pilot_cluster.write_attributes({0x00: 0x02}, manufacturer=0xFC40)
+
+    legrand_wire_pilot_cluster.set_heat_mode.assert_awaited_with(
+        0x02,
+        manufacturer=0xFC40,
+    )
+    legrand_wire_pilot_cluster._write_attributes.assert_awaited_with(
+        [],
+        manufacturer=0xFC40,
+    )
+
+@pytest.mark.parametrize(
+    "attr, value, expected_attr, expected_value",
+    [
+        # Wire pilot mode attribute
+        (0x4000, False, 0x0000, [1, 0]),
+        (0x4000, True, 0x0000, [2, 0]),
+        # Other attributes
+        (0x0001, False, 0x0001, False),
+        (0x0002, True, 0x0002, True),
+    ],
+)
+async def test_legrand_wire_pilot_mode_write_attrs(zigpy_device_from_v2_quirk, attr, value, expected_attr, expected_value):
+    """Test Legrand cable outlet attr writing."""
+
+    device = zigpy_device_from_v2_quirk(f" {LEGRAND}", " Cable outlet")
+    legrand_cluster = device.endpoints[1].legrand_cluster
+    legrand_cluster._write_attributes = mock.AsyncMock()
+
+    await legrand_cluster.write_attributes({ attr: value }, manufacturer=0xFC40)
+
+    expected = foundation.Attribute(expected_attr, foundation.TypeValue())
+    expected_attr_def = legrand_cluster.find_attribute(expected_attr)
+    expected.value.type = foundation.DATA_TYPES.pytype_to_datatype_id(
+        expected_attr_def.type
+    )
+    expected.value.value = expected_attr_def.type(expected_value)
+
+    legrand_cluster._write_attributes.assert_awaited_with(
+        [expected],
+        manufacturer=0xFC40,
+    )
+
+@pytest.mark.parametrize(
+    "attr, value, expected_attr, expected_value",
+    [
+        # Device mode attribute
+        (0x0000, [1, 0], 0x4000, False),
+        (0x0000, [2, 0], 0x4000, True),
+    ],
+)
+async def test_legrand_wire_pilot_mode_update_attr(zigpy_device_from_v2_quirk, attr, value, expected_attr, expected_value):
+    """Test Legrand cable outlet attr update."""
+
+    device = zigpy_device_from_v2_quirk(f" {LEGRAND}", " Cable outlet")
+    legrand_cluster = device.endpoints[1].legrand_cluster
+
+    legrand_cluster_listener = ClusterListener(legrand_cluster)
+
+    legrand_cluster.update_attribute(attr, value)
+
+    assert len(legrand_cluster_listener.attribute_updates) == 2
+    assert legrand_cluster_listener.attribute_updates[0][0] == attr
+    assert legrand_cluster_listener.attribute_updates[0][1] == value
+    assert legrand_cluster_listener.attribute_updates[1][0] == expected_attr
+    assert legrand_cluster_listener.attribute_updates[1][1] == expected_value

zhaquirks/legrand/wire_pilot.py

@@ -0,0 +1,129 @@
+"""Module for Legrand Cable Outlet with pilot wire functionality."""
+
+from zigpy.quirks import CustomCluster
+from zigpy.quirks.v2 import add_to_registry_v2
+import zigpy.types as t
+from zigpy.zcl import ClusterType
+from zigpy.zcl.foundation import (
+    BaseAttributeDefs,
+    BaseCommandDefs,
+    Direction,
+    ZCLAttributeDef,
+    ZCLCommandDef,
+)
+
+from zhaquirks.legrand import LEGRAND, MANUFACTURER_SPECIFIC_CLUSTER_ID
+
+MANUFACTURER_SPECIFIC_CLUSTER_ID_2 = 0xFC40  # 64576
+
+
+class LegrandCluster(CustomCluster):
+    """LegrandCluster."""
+
+    cluster_id = MANUFACTURER_SPECIFIC_CLUSTER_ID
+    name = "LegrandCluster"
+    ep_attribute = "legrand_cluster"
+
+    class AttributeDefs(BaseAttributeDefs):
+        """Attribute definitions."""
+
+        device_mode = ZCLAttributeDef(
+            id=0x0000,
+            type=t.data16,
+            is_manufacturer_specific=True,
+        )
+        led_dark = ZCLAttributeDef(
+            id=0x0001,
+            type=t.Bool,
+            is_manufacturer_specific=True,
+        )
+        led_on = ZCLAttributeDef(
+            id=0x0002,
+            type=t.Bool,
+            is_manufacturer_specific=True,
+        )
+        wire_pilot_mode = ZCLAttributeDef(id=0x4000, type=t.Bool)
+
+    async def write_attributes(self, attributes, manufacturer=None) -> list:
+        """Write attributes to the cluster."""
+
+        attrs = {}
+        for attr, value in attributes.items():
+            attr_def = self.find_attribute(attr)
+            attr_id = attr_def.id
+            if attr_id == LegrandCluster.AttributeDefs.wire_pilot_mode.id:
+                attrs[LegrandCluster.AttributeDefs.device_mode.id] = (
+                    [0x02, 0x00] if value else [0x01, 0x00]
+                )
+            else:
+                attrs[attr] = value
+        return await super().write_attributes(attrs, manufacturer)
+
+    def _update_attribute(self, attrid, value) -> None:
+        super()._update_attribute(attrid, value)
+        if attrid == LegrandCluster.AttributeDefs.device_mode.id:
+            self._update_attribute(
+                LegrandCluster.AttributeDefs.wire_pilot_mode.id, value[0] == 0x02
+            )
+
+
+class HeatMode(t.enum8):
+    """Heat mode."""
+
+    Comfort = 0x00
+    Comfort_minus_1 = 0x01
+    Comfort_minus_2 = 0x02
+    Eco = 0x03
+    Frost_protection = 0x04
+    Off = 0x05
+
+
+class LegrandWirePilotCluster(CustomCluster):
+    """Legrand wire pilot manufacturer-specific cluster."""
+
+    cluster_id = MANUFACTURER_SPECIFIC_CLUSTER_ID_2
+    name = "LegrandWirePilotCluster"
+    ep_attribute = "legrand_wire_pilot_cluster"
+
+    class AttributeDefs(BaseAttributeDefs):
+        """Attribute definitions for LegrandCluster."""
+
+        heat_mode = ZCLAttributeDef(
+            id=0x00,
+            type=HeatMode,
+            is_manufacturer_specific=True,
+        )
+
+    class ServerCommandDefs(BaseCommandDefs):
+        """Server command definitions."""
+
+        set_heat_mode = ZCLCommandDef(
+            id=0x00,
+            schema={"mode": HeatMode},
+            direction=Direction.Client_to_Server,
+            is_manufacturer_specific=True,
+        )
+
+    async def write_attributes(self, attributes, manufacturer=None) -> list:
+        """Write attributes to the cluster."""
+
+        attrs = {}
+        for attr, value in attributes.items():
+            attr_def = self.find_attribute(attr)
+            attr_id = attr_def.id
+            if attr_id == LegrandWirePilotCluster.AttributeDefs.heat_mode.id:
+                await self.set_heat_mode(value, manufacturer=manufacturer)
+        return await super().write_attributes(attrs, manufacturer)
+
+
+(
+    add_to_registry_v2(f" {LEGRAND}", " Cable outlet")
+    .replaces(LegrandCluster)
+    .replaces(LegrandWirePilotCluster)
+    .replaces(LegrandCluster, cluster_type=ClusterType.Client)
+    .switch(
+        attribute_name=LegrandCluster.AttributeDefs.wire_pilot_mode.name,
+        cluster_id=LegrandCluster.cluster_id,
+        translation_key="wire_pilot_mode",
+    )
+)