Fit the Green Power endpoint into two-slot firmware

tests/test_application.py

@@ -8,6 +8,7 @@
 import zigpy.application
 import zigpy.config
 import zigpy.device
+import zigpy.profiles.zll
 from zigpy.types import EUI64, Channels, KeyData
 import zigpy.zdo.types as zdo_t
 
@@ -514,8 +515,6 @@ async def read_param(param_id, index):
 
     app._api.read_parameter = AsyncMock(side_effect=read_param)
     app._api.write_parameter = AsyncMock()
-    app._written_endpoints.add(0x00)
-    app._written_endpoints.add(0x01)
 
     with pytest.raises(ValueError):
         await app.add_endpoint(ENDPOINT.replace(endpoint=1))
@@ -552,6 +551,144 @@ async def read_param(param_id, index):
     )
 
 
+def mock_endpoint_slots(app, slots):
+    """Mock the firmware's endpoint slots, keeping track of writes."""
+
+    async def read_param(param_id, index):
+        assert param_id == deconz_api.NetworkParameter.configure_endpoint
+
+        if index not in slots:
+            raise zigpy_deconz.exception.CommandError(
+                "Unsupported",
+                status=deconz_api.Status.UNSUPPORTED,
+                command=None,
+            )
+
+        return deconz_api.IndexedEndpoint(index=index, descriptor=slots[index])
+
+    async def write_param(param_id, value):
+        assert param_id == deconz_api.NetworkParameter.configure_endpoint
+        assert value.index in slots
+        slots[value.index] = value.descriptor
+
+    app._api.read_parameter = AsyncMock(side_effect=read_param)
+    app._api.write_parameter = AsyncMock(side_effect=write_param)
+
+
+async def zigpy_endpoints(app):
+    """Return the endpoints zigpy registers: HA, ZLL and Green Power."""
+    with patch.object(app, "add_endpoint", AsyncMock()) as add_endpoint:
+        await zigpy.application.ControllerApplication.register_endpoints(app)
+
+    return [call.args[0] for call in add_endpoint.await_args_list]
+
+
+@pytest.mark.parametrize(
+    "initial, expected, num_writes",
+    [
+        # Two slots last configured by zigpy 2.1-2.3: the ZLL endpoint makes room
+        (["ha", "zll"], ["ha", "gp"], 1),
+        # Two slots already holding the HA and GP endpoints are left alone
+        (["ha", "gp"], ["ha", "gp"], 0),
+        # Two slots configured by something else entirely
+        (["other1", "other2"], ["ha", "gp"], 2),
+        # Firmware with more slots keeps the ZLL endpoint
+        (["ha", "zll", "other1", "other2"], ["ha", "zll", "gp", "other2"], 1),
+        (["ha", "zll", "gp"], ["ha", "zll", "gp"], 0),
+    ],
+)
+async def test_register_endpoints(app, initial, expected, num_writes):
+    ha, zll, gp = await zigpy_endpoints(app)
+    named = {
+        "ha": ha,
+        "zll": zll,
+        "gp": gp,
+        "other1": ENDPOINT.replace(endpoint=10),
+        "other2": ENDPOINT.replace(endpoint=11),
+    }
+
+    slots = {index: named[name] for index, name in enumerate(initial)}
+    mock_endpoint_slots(app, slots)
+
+    await app.register_endpoints()
+
+    assert slots == {index: named[name] for index, name in enumerate(expected)}
+    assert app._api.write_parameter.await_count == num_writes
+
+    # Registering again, e.g. when `start_network` is retried, writes nothing
+    app._api.write_parameter.reset_mock()
+    await app.register_endpoints()
+    assert app._api.write_parameter.await_count == 0
+
+
+async def test_register_endpoints_too_many(app):
+    ha, zll, gp = await zigpy_endpoints(app)
+    slots = {0: ha, 1: zll}
+    mock_endpoint_slots(app, slots)
+
+    with patch.dict(
+        app._config,
+        {zigpy.config.CONF_ADDITIONAL_ENDPOINTS: [ENDPOINT.replace(endpoint=10)]},
+    ):
+        with pytest.raises(ValueError):
+            await app.register_endpoints()
+
+    assert slots == {0: ha, 1: zll}
+    app._api.write_parameter.assert_not_called()
+
+
+async def test_register_endpoints_additional_zll(app):
+    """Only zigpy's own ZLL endpoint makes room, a configured one is kept."""
+    ha, zll, gp = await zigpy_endpoints(app)
+    slots = {0: ha, 1: zll}
+    mock_endpoint_slots(app, slots)
+
+    extra = ENDPOINT.replace(endpoint=3, profile=zigpy.profiles.zll.PROFILE_ID)
+
+    with patch.dict(app._config, {zigpy.config.CONF_ADDITIONAL_ENDPOINTS: [extra]}):
+        with pytest.raises(ValueError):
+            await app.register_endpoints()
+
+    assert slots == {0: ha, 1: zll}
+
+
+async def test_register_endpoints_override(app):
+    """A configured endpoint replaces zigpy's endpoint with the same ID."""
+    ha, zll, gp = await zigpy_endpoints(app)
+    slots = {0: ha, 1: zll, 2: gp, 3: ENDPOINT.replace(endpoint=10)}
+    mock_endpoint_slots(app, slots)
+
+    override = ha.replace(device_type=0x1234)
+
+    with patch.dict(app._config, {zigpy.config.CONF_ADDITIONAL_ENDPOINTS: [override]}):
+        await app.register_endpoints()
+
+    assert slots == {0: override, 1: zll, 2: gp, 3: ENDPOINT.replace(endpoint=10)}
+
+
+async def test_add_endpoint_after_register(app):
+    ha, zll, gp = await zigpy_endpoints(app)
+    slots = {0: ha, 1: zll, 2: gp, 3: ENDPOINT.replace(endpoint=10)}
+    mock_endpoint_slots(app, slots)
+
+    await app.register_endpoints()
+    app._api.write_parameter.assert_not_called()
+
+    # A new endpoint goes into the free slot, not one of the registered ones
+    await app.add_endpoint(ENDPOINT.replace(endpoint=20))
+    assert slots == {0: ha, 1: zll, 2: gp, 3: ENDPOINT.replace(endpoint=20)}
+
+    # Replacing an endpoint with the same ID reuses its slot
+    await app.add_endpoint(ENDPOINT.replace(endpoint=20, profile=1234))
+    assert slots == {
+        0: ha,
+        1: zll,
+        2: gp,
+        3: ENDPOINT.replace(endpoint=20, profile=1234),
+    }
+    assert app._api.write_parameter.await_count == 2
+
+
 async def test_reset_network_info(app):
     app.write_network_info = AsyncMock()
     await app.reset_network_info()

zigpy_deconz/zigbee/application.py

@@ -19,6 +19,7 @@
 import zigpy.endpoint
 import zigpy.exceptions
 from zigpy.exceptions import FormationFailure, NetworkNotFormed
+import zigpy.profiles.zll
 import zigpy.state
 import zigpy.types
 import zigpy.util
@@ -57,6 +58,9 @@
 
 CONBEE_III_ENERGY_SCAN_ATTEMPTS = 5
 
+# zigpy registers its ZLL endpoint here, see `ControllerApplication.register_endpoints`
+ZLL_ENDPOINT_ID = 2
+
 
 class ControllerApplication(zigpy.application.ControllerApplication):
     SCHEMA = CONFIG_SCHEMA
@@ -79,7 +83,8 @@ def __init__(self, config: dict[str, Any]):
         self._delayed_neighbor_scan_task = None
         self._reconnect_task = None
 
-        self._written_endpoints = set()
+        self._endpoints: list[zdo_t.SimpleDescriptor] = []
+        self._pending_endpoints: list[zdo_t.SimpleDescriptor] | None = None
 
     async def _watchdog_feed(self):
         if self._api.protocol_version >= PROTO_VER_WATCHDOG and not (
@@ -102,7 +107,7 @@ async def connect(self):
             raise
 
         self._api = api
-        self._written_endpoints.clear()
+        self._endpoints = []
 
     async def disconnect(self):
         if self._delayed_neighbor_scan_task is not None:
@@ -492,56 +497,105 @@ async def _move_network_to_channel(
         await asyncio.sleep(CHANGE_NETWORK_STATE_DELAY)
         await self._change_network_state(NetworkState.CONNECTED)
 
+    async def register_endpoints(self) -> None:
+        """Register all endpoints at once, so they can be fit into the available slots."""
+        self._pending_endpoints = []
+
+        try:
+            await super().register_endpoints()
+            descriptors = self._pending_endpoints
+        finally:
+            self._pending_endpoints = None
+
+        await self._configure_endpoints(descriptors)
+
     async def add_endpoint(self, descriptor: zdo_t.SimpleDescriptor) -> None:
         """Register an endpoint on the device, replacing any with conflicting IDs."""
+        if self._pending_endpoints is not None:
+            self._pending_endpoints.append(descriptor)
+            return
 
-        endpoints = {}
+        await self._configure_endpoints([*self._endpoints, descriptor])
 
-        # Read and count the current endpoints. Some firmwares have three, others four.
+    async def _read_endpoint_slots(self) -> dict[int, zdo_t.SimpleDescriptor]:
+        """Read the endpoint slots of the firmware, they can't be queried otherwise."""
+        slots = {}
+
+        # Most firmwares have two slots, a ConBee II test firmware had four
         for index in range(255 + 1):
             try:
-                current_descriptor = await self._api.read_parameter(
+                current = await self._api.read_parameter(
                     NetworkParameter.configure_endpoint, index
                 )
             except zigpy_deconz.exception.CommandError as ex:
                 assert ex.status == Status.UNSUPPORTED
                 break
             else:
-                endpoints[index] = current_descriptor.descriptor
+                slots[index] = current.descriptor
 
-        LOGGER.debug("Got endpoint slots: %r", endpoints)
+        LOGGER.debug("Got endpoint slots: %r", slots)
 
-        # Don't write endpoints unnecessarily
-        if descriptor in endpoints.values():
-            LOGGER.debug("Endpoint already registered, skipping")
+        return slots
 
-            # Pretend we wrote it
-            self._written_endpoints.add(list(endpoints.values()).index(descriptor))
-            return
+    async def _configure_endpoints(
+        self, descriptors: list[zdo_t.SimpleDescriptor]
+    ) -> None:
+        """Fit the endpoints into the slots, writing only the ones that changed."""
+        slots = await self._read_endpoint_slots()
+
+        # A later endpoint replaces an earlier one with the same ID
+        descriptors = list({ep.endpoint: ep for ep in descriptors}.values())
+
+        # Two-slot firmware cannot host zigpy's ZLL endpoint next to the HA and Green
+        # Power ones. Nothing is sent from it, so drop it like deCONZ and Z2M do.
+        if len(descriptors) > len(slots):
+            fitting = [
+                ep
+                for ep in descriptors
+                if not (
+                    ep.endpoint == ZLL_ENDPOINT_ID
+                    and ep.profile == zigpy.profiles.zll.PROFILE_ID
+                )
+            ]
 
-        # Keep track of the best endpoint descriptor to replace
-        target_index = None
+            if len(fitting) < len(descriptors):
+                LOGGER.debug("Not enough endpoint slots, skipping the ZLL endpoint")
+                descriptors = fitting
 
-        for index, current_descriptor in endpoints.items():
-            # Ignore ones we've already written
-            if index in self._written_endpoints:
-                continue
+        if len(descriptors) > len(slots):
+            raise ValueError(
+                f"Cannot fit {len(descriptors)} endpoints into {len(slots)} endpoint"
+                f" slots: {descriptors!r}"
+            )
 
-            target_index = index
+        assigned: dict[int, zdo_t.SimpleDescriptor] = {}
+        missing: list[zdo_t.SimpleDescriptor] = []
 
-            if current_descriptor.endpoint == descriptor.endpoint:
-                # Prefer to replace the endpoint with the same ID
-                break
+        # Don't write endpoints unnecessarily
+        for descriptor in descriptors:
+            for index, current in slots.items():
+                if index not in assigned and current == descriptor:
+                    assigned[index] = descriptor
+                    break
+            else:
+                missing.append(descriptor)
 
-        if target_index is None:
-            raise ValueError(f"No available endpoint slots exist: {endpoints!r}")
+        free = [index for index in slots if index not in assigned]
 
-        LOGGER.debug("Writing %s to slot %r", descriptor, target_index)
+        for descriptor in missing:
+            # Prefer to replace the endpoint with the same ID
+            index = next(
+                (i for i in free if slots[i].endpoint == descriptor.endpoint), free[0]
+            )
+            free.remove(index)
 
-        await self._api.write_parameter(
-            NetworkParameter.configure_endpoint,
-            IndexedEndpoint(index=target_index, descriptor=descriptor),
-        )
+            LOGGER.debug("Writing %s to slot %r", descriptor, index)
+            await self._api.write_parameter(
+                NetworkParameter.configure_endpoint,
+                IndexedEndpoint(index=index, descriptor=descriptor),
+            )
+
+        self._endpoints = list(descriptors)
 
     async def send_packet(self, packet):
         LOGGER.debug("Sending packet: %r", packet)