Add a unit test

Author: puddly

tests/test_device.py

@@ -18,6 +18,7 @@
 from zigpy.zcl import ClusterType
 from zigpy.zcl.clusters import general
 from zigpy.zcl.clusters.general import Ota, PowerConfiguration
+from zigpy.zcl.clusters.lighting import Color
 from zigpy.zcl.foundation import Status, WriteAttributesResponse
 import zigpy.zdo.types as zdo_t
 
@@ -49,6 +50,8 @@
 from zha.exceptions import ZHAException
 from zha.zigbee.device import (
     ClusterBinding,
+    DeviceEntityAddedEvent,
+    DeviceEntityRemovedEvent,
     DeviceFirmwareInfoUpdatedEvent,
     ZHAEvent,
     get_device_automation_triggers,
@@ -1154,3 +1157,61 @@ async def test_somrig_events(zha_gateway: Gateway) -> None:
             )
         )
     ]
+
+
+async def test_entity_recomputation(zha_gateway: Gateway) -> None:
+    """Test entity recomputation."""
+    zigpy_dev = await zigpy_device_from_json(
+        zha_gateway.application_controller,
+        "tests/data/devices/ikea-of-sweden-tradfri-bulb-gu10-ws-400lm.json",
+    )
+    zha_device = await join_zigpy_device(zha_gateway, zigpy_dev)
+
+    event_listener = mock.Mock()
+    zha_device.on_all_events(event_listener)
+
+    entities1 = set(zha_device.platform_entities.values())
+
+    # We lose track of the color temperature
+    zha_device._zigpy_device.endpoints[1].light_color.add_unsupported_attribute(
+        Color.AttributeDefs.start_up_color_temperature.id
+    )
+    await zha_device.recompute_entities()
+
+    entities2 = set(zha_device.platform_entities.values())
+    assert entities2 - entities1 == set()
+    assert len(entities1 - entities2) == 1
+    assert (
+        list(entities1 - entities2)[0].unique_id
+        == "68:0a:e2:ff:fe:8f:fa:33-1-768-start_up_color_temperature"
+    )
+    assert event_listener.mock_calls == [
+        call(
+            DeviceEntityRemovedEvent(
+                unique_id="68:0a:e2:ff:fe:8f:fa:33-1-768-start_up_color_temperature"
+            )
+        )
+    ]
+
+    event_listener.reset_mock()
+
+    # We add it back
+    zha_device._zigpy_device.endpoints[1].light_color.remove_unsupported_attribute(
+        Color.AttributeDefs.start_up_color_temperature.id
+    )
+    await zha_device.recompute_entities()
+
+    entities3 = set(zha_device.platform_entities.values())
+    assert (
+        list(entities3 - entities2)[0].unique_id
+        == "68:0a:e2:ff:fe:8f:fa:33-1-768-start_up_color_temperature"
+    )
+    assert {e.unique_id for e in entities1} == {e.unique_id for e in entities3}
+
+    assert event_listener.mock_calls == [
+        call(
+            DeviceEntityAddedEvent(
+                unique_id="68:0a:e2:ff:fe:8f:fa:33-1-768-start_up_color_temperature"
+            )
+        )
+    ]