[RFC] Move from centralized discovery to entity-driven discovery

tests/test_button.py

@@ -415,7 +415,7 @@ async def test_cluster_handler_quirks_unnecessary_claiming(
     # Suppress normal endpoint probing, as this will claim the Opple cluster handler
     # already due to it being in the "CLUSTER_HANDLER_ONLY_CLUSTERS" registry.
     # We want to test the handler also gets claimed via quirks v2 attributes init.
-    with patch("zha.application.discovery.EndpointProbe.discover_entities"):
+    with patch("zha.application.discovery.endpoint_discover_entities"):
         zha_device = await join_zigpy_device(zha_gateway, zigpy_device)
     assert isinstance(zha_device.device, CustomDeviceV2)
 

tests/test_cluster_handlers.py

@@ -662,7 +662,7 @@ def test_epch_claim_cluster_handlers(cluster_handler) -> None:
 
 @mock.patch("zha.zigbee.endpoint.Endpoint.add_client_cluster_handlers")
 @mock.patch(
-    "zha.application.discovery.ENDPOINT_PROBE.discover_entities",
+    "zha.application.discovery.endpoint_discover_entities",
     mock.MagicMock(),
 )
 async def test_ep_cluster_handlers_all_cluster_handlers(
@@ -714,7 +714,7 @@ async def test_ep_cluster_handlers_all_cluster_handlers(
 
 @mock.patch("zha.zigbee.endpoint.Endpoint.add_client_cluster_handlers")
 @mock.patch(
-    "zha.application.discovery.ENDPOINT_PROBE.discover_entities",
+    "zha.application.discovery.endpoint_discover_entities",
     mock.MagicMock(),
 )
 async def test_cluster_handler_power_config(
@@ -918,7 +918,7 @@ async def test_zll_device_groups(zha_gateway: Gateway) -> None:
 
 
 @mock.patch(
-    "zha.application.discovery.ENDPOINT_PROBE.discover_entities",
+    "zha.application.discovery.endpoint_discover_entities",
     mock.MagicMock(),
 )
 async def test_cluster_no_ep_attribute(

tests/test_discover.py

@@ -53,12 +53,10 @@
     zigpy_device_from_device_data,
 )
 from zha.application import Platform
-from zha.application.discovery import ENDPOINT_PROBE, EndpointProbe
 from zha.application.gateway import Gateway
 from zha.application.helpers import DeviceOverridesConfiguration
 from zha.application.platforms import PlatformEntity, binary_sensor, sensor
 from zha.application.platforms.number import BaseNumber, NumberMode
-from zha.application.registries import SINGLE_INPUT_CLUSTER_DEVICE_CLASS
 from zha.zigbee.cluster_handlers import ClusterHandler
 from zha.zigbee.endpoint import Endpoint
 
@@ -69,14 +67,14 @@
             return endpoint.identify
 
 
-@mock.patch("zha.application.discovery.EndpointProbe.discover_by_device_type")
-@mock.patch("zha.application.discovery.EndpointProbe.discover_by_cluster_id")
+@mock.patch("zha.application.discovery.endpoint_discover_by_device_type")
+@mock.patch("zha.application.discovery.endpoint_discover_by_cluster_id")
 def test_discover_entities(m1, m2) -> None:
     """Test discover endpoint class method."""
     endpoint = mock.MagicMock()
     endpoint.device.is_coordinator = False
 
-    for _entity in ENDPOINT_PROBE.discover_entities(endpoint, device_overrides={}):
+    for _entity in endpoint_discover_entities(endpoint, device_overrides={}):
         pass
 
     assert m1.call_count == 1
@@ -106,13 +104,11 @@
     entity_cls = mock.MagicMock()
 
     get_entity_mock = mock.MagicMock(return_value=(entity_cls, mock.sentinel.claimed))
     with mock.patch(
         "zha.application.registries.PLATFORM_ENTITIES.get_entity",
         get_entity_mock,
     ):
-        entities = list(
-            ENDPOINT_PROBE.discover_by_device_type(endpoint, device_overrides={})
-        )
+        entities = list(endpoint_discover_by_device_type(endpoint, device_overrides={}))
 
     if hit:
         assert len(entities) == 1
@@ -143,13 +139,13 @@
 
     get_entity_mock = mock.MagicMock(return_value=(entity_cls, mock.sentinel.claimed))
     with (
         mock.patch(
             "zha.application.registries.PLATFORM_ENTITIES.get_entity",
             get_entity_mock,
         ),
     ):
         entities = list(
-            ENDPOINT_PROBE.discover_by_device_type(
+            endpoint_discover_by_device_type(
                 endpoint,
                 device_overrides={
                     "00:11:22:33:44:55:66:77-1": DeviceOverridesConfiguration(
@@ -184,11 +180,11 @@
     entity_cls = mock.MagicMock()
     get_entity_mock = mock.MagicMock(return_value=(entity_cls, mock.sentinel.claimed))
     cluster_handler_mock = mock.MagicMock(spec_set=ClusterHandler)
     with mock.patch(
         "zha.application.registries.PLATFORM_ENTITIES.get_entity",
         get_entity_mock,
     ):
-        for _entity in ENDPOINT_PROBE.probe_single_cluster(
+        for _entity in probe_single_cluster(
             Platform.SWITCH, cluster_handler_mock, endpoint
         ):
             pass
@@ -214,60 +210,6 @@
     return cluster_handler
 
 
-def test_single_input_cluster_device_class_by_cluster_class() -> None:
-    """Test SINGLE_INPUT_CLUSTER_DEVICE_CLASS matching by cluster id or class."""
-
-    class QuirkedIAS(zigpy.quirks.CustomCluster, zigpy.zcl.clusters.security.IasZone):
-        """Quirked IAS Zone cluster."""
-
-    class _Analog(zigpy.quirks.CustomCluster, zigpy.zcl.clusters.general.AnalogInput):
-        pass
-
-    door_ch = _ch_mock(zigpy.zcl.clusters.closures.DoorLock)
-    cover_ch = _ch_mock(zigpy.zcl.clusters.closures.WindowCovering)
-    multistate_ch = _ch_mock(zigpy.zcl.clusters.general.MultistateInput)
-    ias_ch = _ch_mock(QuirkedIAS)
-    analog_ch = _ch_mock(_Analog)
-
-    endpoint = mock.MagicMock(spec_set=Endpoint)
-    endpoint.unclaimed_cluster_handlers.return_value = [
-        door_ch,
-        cover_ch,
-        multistate_ch,
-        ias_ch,
-        analog_ch,
-    ]
-
-    with (
-        mock.patch.dict(
-            SINGLE_INPUT_CLUSTER_DEVICE_CLASS,
-            {
-                zigpy.zcl.clusters.closures.DoorLock.cluster_id: Platform.LOCK,
-                zigpy.zcl.clusters.closures.WindowCovering.cluster_id: Platform.COVER,
-                zigpy.zcl.clusters.general.AnalogInput.cluster_id: Platform.SENSOR,
-                zigpy.zcl.clusters.general.MultistateInput.cluster_id: Platform.SENSOR,
-                zigpy.zcl.clusters.security.IasZone.cluster_id: Platform.BINARY_SENSOR,
-            },
-            clear=True,
-        ),
-        mock.patch(
-            "zha.application.discovery.EndpointProbe.probe_single_cluster",
-            new=mock.MagicMock(),
-        ) as probe_mock,
-    ):
-        for _entity in EndpointProbe().discover_by_cluster_id(endpoint):
-            pass
-
-        assert probe_mock.call_count == len(endpoint.unclaimed_cluster_handlers())
-        assert [m for m in probe_mock.mock_calls if m != call().__iter__()] == [
-            call(Platform.LOCK, door_ch, endpoint),
-            call(Platform.COVER, cover_ch, endpoint),
-            call(Platform.SENSOR, multistate_ch, endpoint),
-            call(Platform.BINARY_SENSOR, ias_ch, endpoint),
-            call(Platform.SENSOR, analog_ch, endpoint),
-        ]
-
-
 @pytest.mark.parametrize("override", [None, "switch"])
 async def test_device_override(
     zha_gateway: Gateway,
@@ -302,7 +244,7 @@
     await zha_gateway.async_block_till_done()
     zha_device = zha_gateway.get_device(zigpy_device.ieee)
 
     get_entity(zha_device, platform=Platform.SWITCH if override else Platform.LIGHT)
 
 
 async def test_quirks_v2_entity_discovery(
@@ -825,7 +767,7 @@
                     f"Duplicate unique_id {unique_id} found in entities: {entities}"
                 )
             else:
                 warnings.warn(
                     f"Unique IDs are unique only with platform prefix: {dict(zip(prefixed_unique_ids, entities))}"
                 )
 
@@ -848,7 +790,7 @@
         loaded_device_data = json.loads(
             json.dumps(zha_device.get_diagnostics_json(), cls=ZhaJsonEncoder)
         )
         assert loaded_device_data == device_data
 
         # Assert identify called on join for devices that support it
         cluster_identify = _get_identify_cluster(zha_device.device)

tests/test_sensor.py

@@ -1604,7 +1604,7 @@ async def test_cluster_handler_quirks_attribute_reporting(zha_gateway: Gateway)
     # Suppress normal endpoint probing, as this will claim the Opple cluster handler
     # already due to it being in the "CLUSTER_HANDLER_ONLY_CLUSTERS" registry.
     # We want to test the handler also gets claimed via quirks v2 reporting config.
-    with patch("zha.application.discovery.EndpointProbe.discover_entities"):
+    with patch("zha.application.discovery.endpoint_discover_entities"):
         zha_device, cluster = await zigpy_device_aqara_sensor_v2_mock(zha_gateway)
     assert isinstance(zha_device.device, CustomDeviceV2)
 
@@ -1687,7 +1687,7 @@ async def test_cluster_handler_quirks_attribute_reading(zha_gateway: Gateway) ->
     # Suppress normal endpoint probing, as this will claim the Opple cluster handler
     # already due to it being in the "CLUSTER_HANDLER_ONLY_CLUSTERS" registry.
     # We want to test the handler also gets claimed via quirks v2 attributes init.
-    with patch("zha.application.discovery.EndpointProbe.discover_entities"):
+    with patch("zha.application.discovery.endpoint_discover_entities"):
         zha_device = await join_zigpy_device(zha_gateway, zigpy_device)
     assert isinstance(zha_device.device, CustomDeviceV2)