Implement quirks v2 unique_id_suffix

tests/test_button.py

@@ -14,8 +14,8 @@
 from zhaquirks.tuya.ts0601_valve import ParksideTuyaValveManufCluster
 from zigpy.exceptions import ZigbeeException
 from zigpy.profiles import zha
-from zigpy.quirks import CustomCluster, CustomDevice
-from zigpy.quirks.v2 import QuirkBuilder
+from zigpy.quirks import CustomCluster, CustomDevice, DeviceRegistry
+from zigpy.quirks.v2 import CustomDeviceV2, QuirkBuilder
 import zigpy.types as t
 from zigpy.zcl.clusters import general, security
 from zigpy.zcl.clusters.manufacturer_specific import ManufacturerSpecificCluster
@@ -30,6 +30,7 @@
     get_entity,
     join_zigpy_device,
     mock_coro,
+    patch_cluster_for_testing,
     update_attribute_cache,
 )
 from zha.application import Platform
@@ -184,30 +185,10 @@ class ServerCommandDefs(zcl_f.BaseCommandDefs):
         )
 
 
-(
-    QuirkBuilder("Fake_Model", "Fake_Manufacturer")
-    .replaces(FakeManufacturerCluster)
-    .command_button(
-        FakeManufacturerCluster.ServerCommandDefs.self_test.name,
-        FakeManufacturerCluster.cluster_id,
-        command_args=(5,),
-        translation_key="self_test",
-        fallback_name="Self test",
-    )
-    .write_attr_button(
-        FakeManufacturerCluster.AttributeDefs.feed.name,
-        2,
-        FakeManufacturerCluster.cluster_id,
-        translation_key="feed",
-        fallback_name="Feed",
-    )
-    .add_to_registry()
-)
-
-
 async def custom_button_device(zha_gateway: Gateway):
     """Button device fixture for quirks button tests."""
 
+    registry = DeviceRegistry()
     zigpy_device = create_mock_zigpy_device(
         zha_gateway,
         {
@@ -225,6 +206,32 @@ async def custom_button_device(zha_gateway: Gateway):
         model="Fake_Manufacturer",
     )
 
+    (
+        QuirkBuilder("Fake_Model", "Fake_Manufacturer", registry=registry)
+        .replaces(FakeManufacturerCluster)
+        .command_button(
+            FakeManufacturerCluster.ServerCommandDefs.self_test.name,
+            FakeManufacturerCluster.cluster_id,
+            command_args=(5,),
+            translation_key="self_test",
+            fallback_name="Self test",
+        )
+        .write_attr_button(
+            FakeManufacturerCluster.AttributeDefs.feed.name,
+            2,
+            FakeManufacturerCluster.cluster_id,
+            translation_key="feed",
+            fallback_name="Feed",
+        )
+        .add_to_registry()
+    )
+
+    zigpy_device = registry.get_device(zigpy_device)
+
+    assert isinstance(zigpy_device, CustomDeviceV2)
+    # XXX: this should be handled automatically, patch quirks added cluster
+    patch_cluster_for_testing(zigpy_device.endpoints[1].mfg_identify)
+
     zigpy_device.endpoints[1].mfg_identify.PLUGGED_ATTR_READS = {
         FakeManufacturerCluster.AttributeDefs.feed.name: 0,
     }
@@ -276,3 +283,66 @@ async def test_quirks_write_attr_button(
         ]
 
     assert cluster.get(cluster.AttributeDefs.feed.name) == 2
+
+
+async def test_quirks_write_attr_buttons_uid(zha_gateway: Gateway) -> None:
+    """Test multiple buttons created with different unique id suffixes."""
+
+    registry = DeviceRegistry()
+    zigpy_dev = create_mock_zigpy_device(
+        zha_gateway,
+        {
+            1: {
+                SIG_EP_INPUT: [
+                    general.Basic.cluster_id,
+                    FakeManufacturerCluster.cluster_id,
+                ],
+                SIG_EP_OUTPUT: [],
+                SIG_EP_TYPE: zha.DeviceType.REMOTE_CONTROL,
+                SIG_EP_PROFILE: zha.PROFILE_ID,
+            }
+        },
+        manufacturer="Fake_Model",
+        model="Fake_Manufacturer",
+    )
+
+    (
+        QuirkBuilder("Fake_Model", "Fake_Manufacturer", registry=registry)
+        .replaces(FakeManufacturerCluster)
+        .write_attr_button(
+            FakeManufacturerCluster.AttributeDefs.feed.name,
+            1,
+            FakeManufacturerCluster.cluster_id,
+            unique_id_suffix="btn_1",
+            translation_key="btn_1",
+            fallback_name="Button 1",
+        )
+        .write_attr_button(
+            FakeManufacturerCluster.AttributeDefs.feed.name,
+            2,
+            FakeManufacturerCluster.cluster_id,
+            unique_id_suffix="btn_2",
+            translation_key="btn_2",
+            fallback_name="Button 2",
+        )
+        .add_to_registry()
+    )
+
+    zigpy_device_ = registry.get_device(zigpy_dev)
+
+    assert isinstance(zigpy_device_, CustomDeviceV2)
+    zha_device = await join_zigpy_device(zha_gateway, zigpy_device_)
+
+    entity_btn_1 = get_entity(zha_device, platform=Platform.BUTTON, qualifier="btn_1")
+    entity_btn_2 = get_entity(zha_device, platform=Platform.BUTTON, qualifier="btn_2")
+
+    # check both entities are created and have a different unique id suffix
+    assert isinstance(entity_btn_1, WriteAttributeButton)
+    assert entity_btn_1.translation_key == "btn_1"
+    assert entity_btn_1._unique_id_suffix == "btn_1"
+    assert entity_btn_1._attribute_value == 1
+
+    assert isinstance(entity_btn_2, WriteAttributeButton)
+    assert entity_btn_2.translation_key == "btn_2"
+    assert entity_btn_2._unique_id_suffix == "btn_2"
+    assert entity_btn_2._attribute_value == 2

zha/application/platforms/__init__.py

@@ -350,7 +350,9 @@ def _init_from_quirks_metadata(self, entity_metadata: EntityMetadata) -> None:
         if entity_metadata.translation_key:
             self._attr_translation_key = entity_metadata.translation_key
 
-        if has_attribute_name:
+        if unique_id_suffix := entity_metadata.unique_id_suffix:
+            self._unique_id_suffix = unique_id_suffix
+        elif has_attribute_name:
             self._unique_id_suffix = entity_metadata.attribute_name
         elif has_command_name:
             self._unique_id_suffix = entity_metadata.command_name