Entity recomputation and add/remove at runtime

tests/test_device.py

@@ -23,6 +23,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.clusters.measurement import CarbonDioxideConcentration
 from zigpy.zcl.foundation import Status, WriteAttributesResponse
 from zigpy.zcl.helpers import ReportingConfig
@@ -56,6 +57,8 @@
 from zha.exceptions import ZHAException
 from zha.zigbee.device import (
     ClusterBinding,
+    DeviceEntityAddedEvent,
+    DeviceEntityRemovedEvent,
     DeviceFirmwareInfoUpdatedEvent,
     ZHAEvent,
     get_device_automation_triggers,
@@ -1360,6 +1363,99 @@ async def test_device_on_remove_pending_entity_failure(
     assert "Pending entity removal failed" in caplog.text
 
 
+async def test_initial_entity_discovery_does_not_emit_events(
+    zha_gateway: Gateway,
+) -> None:
+    """Test that first device initialization does not emit entity events."""
+    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)
+
+    # Reset to pre-initialization state
+    for entity in list(zha_device.platform_entities.values()):
+        await zha_device._remove_entity(entity, emit_event=False)
+    zha_device._initialized = False
+
+    event_listener = mock.Mock()
+    zha_device.on_event(DeviceEntityAddedEvent.event_type, event_listener)
+
+    # First initialization: entities are discovered but no events should fire
+    await zha_device.async_initialize(from_cache=True)
+
+    assert len(zha_device.platform_entities) > 0
+    assert event_listener.call_count == 0
+
+
+async def test_reinitialize_emits_events_for_new_entities(
+    zha_gateway: Gateway,
+) -> None:
+    """Test that re-initializing a device emits events for new entities."""
+    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)
+
+    unique_id = "68:0a:e2:ff:fe:8f:fa:33-1-768-start_up_color_temperature"
+    entity = zha_device.get_platform_entity(Platform.NUMBER, unique_id)
+    await zha_device._remove_entity(entity, emit_event=False)
+
+    event_listener = mock.Mock()
+    zha_device.on_event(DeviceEntityAddedEvent.event_type, event_listener)
+
+    # Re-initialize with existing entities: new entity should emit an event
+    await zha_device.async_initialize(from_cache=True)
+
+    assert event_listener.call_count == 1
+    assert event_listener.call_args[0][0] == DeviceEntityAddedEvent(
+        platform=Platform.NUMBER,
+        unique_id=unique_id,
+    )
+
+
+async def test_reinitialize_after_on_remove_emits_events(
+    zha_gateway: Gateway,
+) -> None:
+    """Test that re-init after on_remove (all entities cleared) still emits events."""
+    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)
+    entity_count = len(zha_device.platform_entities)
+
+    # Simulate a full removal, clearing all entities
+    await zha_device.on_remove()
+    assert len(zha_device.platform_entities) == 0
+
+    event_listener = mock.Mock()
+    zha_device.on_event(DeviceEntityAddedEvent.event_type, event_listener)
+
+    await zha_device.async_initialize(from_cache=True)
+
+    assert len(zha_device.platform_entities) == entity_count
+    assert event_listener.call_count == entity_count
+
+
+async def test_remove_entity_no_event(zha_gateway: Gateway) -> None:
+    """Test that _remove_entity with emit_event=False does not emit."""
+    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_event(DeviceEntityRemovedEvent.event_type, event_listener)
+
+    existing_entity = next(iter(zha_device.platform_entities.values()))
+    await zha_device._remove_entity(existing_entity, emit_event=False)
+
+    assert event_listener.call_count == 0
+
+
 async def test_initialize_endpoint_failure(zha_gateway: Gateway) -> None:
     """Test that a failing endpoint doesn't prevent device initialization."""
     zigpy_dev = await zigpy_device_from_json(
@@ -1376,3 +1472,93 @@ async def test_initialize_endpoint_failure(zha_gateway: Gateway) -> None:
         await zha_device.async_initialize(from_cache=True)
 
     assert mock_async_initialize.call_count == 1
+
+
+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(
+                platform=Platform.NUMBER,
+                unique_id="68:0a:e2:ff:fe:8f:fa:33-1-768-start_up_color_temperature",
+                remove=True,
+            )
+        )
+    ]
+
+    event_listener.reset_mock()
+
+    # We add it back by writing a value, which clears the unsupported flag
+    zha_device._zigpy_device.endpoints[1].light_color.update_attribute(
+        Color.AttributeDefs.start_up_color_temperature.id, 250
+    )
+    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(
+                platform=Platform.NUMBER,
+                unique_id="68:0a:e2:ff:fe:8f:fa:33-1-768-start_up_color_temperature",
+            )
+        )
+    ]
+
+
+async def test_add_entity_duplicate(zha_gateway: Gateway) -> None:
+    """Test that adding a duplicate entity raises an error."""
+    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)
+
+    existing_entity = next(iter(zha_device.platform_entities.values()))
+
+    with pytest.raises(ValueError, match="unique ID already taken"):
+        zha_device._add_entity(existing_entity)
+
+
+async def test_remove_entity_nonexistent(zha_gateway: Gateway) -> None:
+    """Test that removing a nonexistent entity raises an error."""
+    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)
+
+    existing_entity = next(iter(zha_device.platform_entities.values()))
+    await zha_device._remove_entity(existing_entity)
+
+    with pytest.raises(ValueError, match="unique ID not found"):
+        await zha_device._remove_entity(existing_entity)

tests/test_discover.py

@@ -714,6 +714,9 @@ async def test_devices_from_files(
             await zha_gateway.async_block_till_done(wait_background_tasks=True)
             assert zha_device is not None
 
+        # Ensure entity recomputation is idempotent
+        await zha_device.recompute_entities()
+
         unique_id_collisions = defaultdict(list)
         for entity in zha_device.platform_entities.values():
             unique_id_collisions[entity.unique_id].append(entity)
@@ -748,8 +751,6 @@ async def test_devices_from_files(
 
                 unique_id_migrations[key] = entity
 
-        await zha_device.on_remove()
-
         # XXX: We re-serialize the JSON because integer enum types are converted when
         # serializing but will not compare properly otherwise
         loaded_device_data = json.loads(
@@ -778,6 +779,8 @@ async def test_devices_from_files(
                 )
             ]
 
+        await zha_device.on_remove()
+
 
 async def test_get_diagnostics_json_repeated_calls(zha_gateway: Gateway) -> None:
     """Test that calling get_diagnostics_json twice produces the same result."""

zha/application/const.py

@@ -193,6 +193,8 @@ def pretty_name(self) -> str:
 ZHA_CLUSTER_HANDLER_READS_PER_REQ = 5
 ZHA_EVENT = "zha_event"
 ZHA_DEVICE_UPDATED_EVENT = "zha_device_updated_event"
+ZHA_DEVICE_ENTITY_ADDED_EVENT = "zha_device_entity_added_event"
+ZHA_DEVICE_ENTITY_REMOVED_EVENT = "zha_device_entity_removed_event"
 ZHA_GW_MSG = "zha_gateway_message"
 ZHA_GW_MSG_DEVICE_FULL_INIT = "device_fully_initialized"
 ZHA_GW_MSG_DEVICE_INFO = "device_info"

zha/application/platforms/button/__init__.py

@@ -147,7 +147,7 @@ class IdentifyButton(Button):
     def is_supported_in_list(self, entities: list[BaseEntity]) -> bool:
         """Check if this button is supported given the list of entities."""
         cls = type(self)
-        return not any(type(entity) is cls for entity in entities)
+        return not any(type(entity) is cls for entity in entities if entity is not self)
 
 
 class WriteAttributeButton(BaseButton):

zha/application/platforms/sensor/__init__.py

@@ -2039,7 +2039,7 @@ def _is_supported(self) -> bool:
     def is_supported_in_list(self, entities: list[BaseEntity]) -> bool:
         """Check if the sensor is supported given the list of entities."""
         cls = type(self)
-        return not any(type(entity) is cls for entity in entities)
+        return not any(type(entity) is cls for entity in entities if entity is not self)
 
     @property
     def state(self) -> dict: