Clean up when removing a device from the network (#91)

Author: dmulcahey

tests/conftest.py

@@ -385,6 +385,7 @@ def device_joined(
     """Return a newly joined ZHAWS device."""
 
     async def _zha_device(zigpy_dev: zigpy.device.Device) -> Device:
+        zha_gateway.application_controller.devices[zigpy_dev.ieee] = zigpy_dev
         await zha_gateway.async_device_initialized(zigpy_dev)
         await zha_gateway.async_block_till_done()
 

tests/test_gateway.py

@@ -305,6 +305,59 @@ async def test_gateway_create_group_with_id_without_id(
     assert zha_group3.group_id == 0x0003
 
 
+async def test_remove_device_cleans_up_group_membership(
+    zha_gateway: Gateway,
+    device_light_1,  # pylint: disable=redefined-outer-name
+    device_light_2,  # pylint: disable=redefined-outer-name
+    coordinator,  # pylint: disable=redefined-outer-name
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test group membership cleanup when removing a device in a group."""
+
+    zha_gateway.coordinator_zha_device = coordinator
+    coordinator._zha_gateway = zha_gateway
+    device_light_1._zha_gateway = zha_gateway
+    device_light_2._zha_gateway = zha_gateway
+
+    member_ieee_addresses = [device_light_1.ieee, device_light_2.ieee]
+    members = [
+        GroupMemberReference(ieee=device_light_1.ieee, endpoint_id=1),
+        GroupMemberReference(ieee=device_light_2.ieee, endpoint_id=1),
+    ]
+
+    # test creating a group with 2 members
+    zha_group: Group = await zha_gateway.async_create_zigpy_group("Test Group", members)
+    await zha_gateway.async_block_till_done()
+
+    assert zha_group is not None
+    assert len(zha_group.members) == 2
+    for member in zha_group.members:
+        assert member.device.ieee in member_ieee_addresses
+        assert member.group == zha_group
+        assert member.endpoint is not None
+
+    await zha_gateway.async_remove_device(coordinator.ieee)
+    await zha_gateway.async_block_till_done()
+    assert len(zha_group.members) == 2
+    assert (
+        f"Removing the active coordinator ({str(coordinator.ieee)}) is not allowed"
+        in caplog.text
+    )
+
+    non_existent_ieee = zigpy.types.EUI64.convert("01:2d:6f:70:7a:40:79:e8")
+    await zha_gateway.async_remove_device(non_existent_ieee)
+    await zha_gateway.async_block_till_done()
+    assert len(zha_group.members) == 2
+    assert f"Device: {str(non_existent_ieee)} could not be found" in caplog.text
+
+    await zha_gateway.async_remove_device(device_light_1.ieee)
+    await zha_gateway.async_block_till_done()
+
+    assert len(zha_group.members) == 1
+    assert zha_group.members[0].device.ieee == device_light_2.ieee
+    assert device_light_1.ieee not in zha_gateway.devices
+
+
 @patch(
     "zha.application.gateway.Gateway.load_devices",
     MagicMock(),

zha/application/gateway.py

@@ -670,6 +670,21 @@ async def async_create_zigpy_group(
                 await asyncio.gather(*tasks)
         return self.groups.get(group_id)
 
+    async def async_remove_device(self, ieee: EUI64) -> None:
+        """Remove a device from ZHA."""
+        if not (device := self.devices.get(ieee)):
+            _LOGGER.debug("Device: %s could not be found", ieee)
+            return
+        if device.is_active_coordinator:
+            _LOGGER.info("Removing the active coordinator (%s) is not allowed", ieee)
+            return
+        for group_id, group in self.groups.items():
+            for member_ieee_endpoint_id in list(group.zigpy_group.members.keys()):
+                if member_ieee_endpoint_id[0] == ieee:
+                    await device.async_remove_from_group(group_id)
+
+        await self.application_controller.remove(ieee)
+
     async def async_remove_zigpy_group(self, group_id: int) -> None:
         """Remove a Zigbee group from Zigpy."""
         if not (group := self.groups.get(group_id)):