Explicit entity object discovery, lifetimes, and registration (#365)

* Make entity lifecycle more explicit

* Fix tests

* WIP: Clean up a bit more

* WIP: Add `_remove_callbacks`

* WIP: Centralize entity registration

* Move over a few more `__init__` event handlers

* WIP: Replace `create_platform_entity` with `is_supported`

* WIP: Remove device initialization checks and actually use `is_supported`

* Initialize the device only after entities are created

* WIP

* Explicitly register cluster handlers

* Test: introduce `Entity.recompute_capabilities()`

* Test: migrate a few more entities to `recompute_capabilities`

* Only check if an entity is supported after cluster handler initialization

* One more `recompute_capabilities` for sensor

* Test: add `Entity._always_supported` to properly deal with v2 quirks

* Fix remaining failing non-discovery tests

* Fix cluster handler edge case not calling `on_add()` (thanks @dmulcahey!)

* Clean up definition format

* Fix typing

* Fix pre-commit issues

* Fix remaining tests

* Rename `_maybe_create_device_entities` to `_maybe_create_entities`

* Fix last failing unit test

* Give `get_entity` a strict mode

* Fix remaining tests

* Drop strict mode by fixing violations

* "Fix" final unit test

* Implement `endpoint.on_remove`

* Move entity creation onto the device object

* Make initialization and configuration logic more explicit

* Retain the old behavior of implicit initialization during rejoin

* Remove `single_device_matches` in favor of explicit entity checks

WIP

* Clear out unsupported entities in reverse order

Author: puddly

tests/common.py

@@ -231,6 +231,8 @@ def get_entity(
     qualifier_func: Callable[[BaseEntity], bool] = lambda e: True,
 ) -> PlatformEntity:
     """Get the first entity of the specified platform on the given device."""
+    results = []
+
     for entity in device.platform_entities.values():
         if platform != entity.PLATFORM:
             continue
@@ -247,11 +249,19 @@ def get_entity(
         if not qualifier_func(entity):
             continue
 
-        return entity
+        results.append(entity)
 
-    raise KeyError(
-        f"No {entity_type} entity found for platform {platform!r} on device {device}: {device.platform_entities}"
-    )
+    if len(results) == 0:
+        raise KeyError(
+            f"No {entity_type} entity found for platform {platform!r} on device {device}: {device.platform_entities}"
+        )
+
+    if len(results) != 1:
+        raise KeyError(
+            f"Multiple {entity_type} entities found for platform {platform!r} on device {device}: {results}"
+        )
+
+    return results[0]
 
 
 async def group_entity_availability_test(

tests/test_button.py

@@ -36,7 +36,11 @@
 from zha.application import Platform
 from zha.application.gateway import Gateway
 from zha.application.platforms import EntityCategory, PlatformEntity
-from zha.application.platforms.button import Button, WriteAttributeButton
+from zha.application.platforms.button import (
+    Button,
+    FrostLockResetButton,
+    WriteAttributeButton,
+)
 from zha.application.platforms.button.const import ButtonDeviceClass
 from zha.exceptions import ZHAException
 from zha.zigbee.device import Device
@@ -136,9 +140,11 @@ async def test_frost_unlock(
     cluster = zigpy_device.endpoints[1].tuya_manufacturer
     assert cluster is not None
     entity: PlatformEntity = get_entity(
-        zha_device, platform=Platform.BUTTON, entity_type=WriteAttributeButton
+        zha_device,
+        platform=Platform.BUTTON,
+        entity_type=FrostLockResetButton,
     )
-    assert isinstance(entity, WriteAttributeButton)
+    assert isinstance(entity, FrostLockResetButton)
 
     assert entity._attr_device_class == ButtonDeviceClass.RESTART
     assert entity._attr_entity_category == EntityCategory.CONFIG
@@ -246,7 +252,9 @@ async def test_quirks_command_button(
     """Test ZHA button platform."""
     zha_device, cluster = await custom_button_device(zha_gateway)
     assert cluster is not None
-    entity: PlatformEntity = get_entity(zha_device, platform=Platform.BUTTON)
+    entity: PlatformEntity = get_entity(
+        zha_device, platform=Platform.BUTTON, entity_type=Button
+    )
 
     with patch(
         "zigpy.zcl.Cluster.request",

tests/test_climate.py

@@ -290,51 +290,61 @@ async def test_climate_hvac_action_running_state(
         dev_climate_sinope, platform=Platform.SENSOR, entity_type=SinopeHVACAction
     )
 
-    subscriber = MagicMock()
-    entity.on_event(STATE_CHANGED, subscriber)
-    sensor_entity.on_event(STATE_CHANGED, subscriber)
+    subscriber1 = MagicMock()
+    subscriber2 = MagicMock()
+    entity.on_event(STATE_CHANGED, subscriber1)
+    sensor_entity.on_event(STATE_CHANGED, subscriber2)
 
     assert entity.state["hvac_action"] == "off"
     assert sensor_entity.state["state"] == "off"
 
     await send_attributes_report(
         zha_gateway, thrm_cluster, {0x001E: Thermostat.RunningMode.Off}
     )
+    await zha_gateway.async_block_till_done(wait_background_tasks=True)
     assert entity.state["hvac_action"] == "off"
     assert sensor_entity.state["state"] == "off"
+    assert len(subscriber1.mock_calls) == len(subscriber2.mock_calls) == 0
 
     await send_attributes_report(
         zha_gateway, thrm_cluster, {0x001C: Thermostat.SystemMode.Auto}
     )
+    await zha_gateway.async_block_till_done(wait_background_tasks=True)
     assert entity.state["hvac_action"] == "idle"
     assert sensor_entity.state["state"] == "idle"
+    assert len(subscriber1.mock_calls) == len(subscriber2.mock_calls) == 1
 
     await send_attributes_report(
         zha_gateway, thrm_cluster, {0x001E: Thermostat.RunningMode.Cool}
     )
+    await zha_gateway.async_block_till_done(wait_background_tasks=True)
     assert entity.state["hvac_action"] == "cooling"
     assert sensor_entity.state["state"] == "cooling"
+    assert len(subscriber1.mock_calls) == len(subscriber2.mock_calls) == 2
 
     await send_attributes_report(
         zha_gateway, thrm_cluster, {0x001E: Thermostat.RunningMode.Heat}
     )
+    await zha_gateway.async_block_till_done(wait_background_tasks=True)
     assert entity.state["hvac_action"] == "heating"
     assert sensor_entity.state["state"] == "heating"
+    assert len(subscriber1.mock_calls) == len(subscriber2.mock_calls) == 3
 
     await send_attributes_report(
         zha_gateway, thrm_cluster, {0x001E: Thermostat.RunningMode.Off}
     )
+    await zha_gateway.async_block_till_done(wait_background_tasks=True)
     assert entity.state["hvac_action"] == "idle"
     assert sensor_entity.state["state"] == "idle"
+    assert len(subscriber1.mock_calls) == len(subscriber2.mock_calls) == 4
 
     await send_attributes_report(
         zha_gateway, thrm_cluster, {0x0029: Thermostat.RunningState.Fan_State_On}
     )
+    await zha_gateway.async_block_till_done(wait_background_tasks=True)
     assert entity.state["hvac_action"] == "fan"
     assert sensor_entity.state["state"] == "fan"
-
-    # Both entities are updated!
-    assert len(subscriber.mock_calls) == 2 * 6
+    assert len(subscriber1.mock_calls) == len(subscriber2.mock_calls) == 5
 
 
 async def test_sinope_time(