Add guards to prevent shutdown from failing (#542)

* Add guards to prevent shutdown from failing

* Add unit tests

Author: puddly

tests/test_device.py

@@ -1212,3 +1212,70 @@ async def test_symfonisk_events(
             )
         )
     ]
+
+
+async def test_device_on_remove_callback_failure(
+    zha_gateway: Gateway,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test that device.on_remove continues when callback fails."""
+    zigpy_dev = await zigpy_device_from_json(
+        zha_gateway.application_controller,
+        "tests/data/devices/philips-sml001.json",
+    )
+    zha_device = await join_zigpy_device(zha_gateway, zigpy_dev)
+
+    failing_callback = mock.Mock(side_effect=Exception("Callback failed"))
+    zha_device._on_remove_callbacks.append(failing_callback)
+
+    await zha_device.on_remove()
+
+    assert failing_callback.call_count == 1
+    assert "Failed to execute on_remove callback" in caplog.text
+    assert "Callback failed" in caplog.text
+
+
+async def test_device_on_remove_platform_entity_failure(
+    zha_gateway: Gateway,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test that device.on_remove continues when platform entity removal fails."""
+    zigpy_dev = await zigpy_device_from_json(
+        zha_gateway.application_controller,
+        "tests/data/devices/philips-sml001.json",
+    )
+    zha_device = await join_zigpy_device(zha_gateway, zigpy_dev)
+
+    switch_entity = get_entity(zha_device, platform=Platform.SWITCH)
+    with patch.object(
+        switch_entity, "on_remove", side_effect=Exception("Entity removal failed")
+    ):
+        await zha_device.on_remove()
+
+    assert "Failed to remove platform entity" in caplog.text
+    assert "Entity removal failed" in caplog.text
+
+
+async def test_device_on_remove_pending_entity_failure(
+    zha_gateway: Gateway,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test that device.on_remove continues when pending entity removal fails."""
+    zigpy_dev = await zigpy_device_from_json(
+        zha_gateway.application_controller,
+        "tests/data/devices/philips-sml001.json",
+    )
+    zha_device = await join_zigpy_device(zha_gateway, zigpy_dev)
+
+    update_entity = get_entity(zha_device, platform=Platform.UPDATE)
+    zha_device._pending_entities.append(update_entity)
+
+    with patch.object(
+        update_entity,
+        "on_remove",
+        side_effect=Exception("Pending entity removal failed"),
+    ):
+        await zha_device.on_remove()
+
+    assert "Failed to remove pending entity" in caplog.text
+    assert "Pending entity removal failed" in caplog.text

tests/test_gateway.py

@@ -21,6 +21,7 @@
     get_entity,
     get_group_entity,
     join_zigpy_device,
+    zigpy_device_from_json,
 )
 from zha.application import Platform
 from zha.application.const import (
@@ -882,3 +883,86 @@ async def test_gateway_energy_scan(zha_gateway: Gateway) -> None:
         assert mock_scan.mock_calls == [
             call(channels=channels, duration_exp=4, count=1)
         ]
+
+
+async def test_gateway_shutdown_device_on_remove_failure(
+    zha_gateway: Gateway,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test that gateway shutdown continues when device.on_remove fails."""
+    zigpy_dev = await zigpy_device_from_json(
+        zha_gateway.application_controller,
+        "tests/data/devices/centralite-3320-l.json",
+    )
+    zha_device = await join_zigpy_device(zha_gateway, zigpy_dev)
+
+    with patch.object(
+        zha_device, "on_remove", side_effect=Exception("Device removal failed")
+    ):
+        await zha_gateway.shutdown()
+        await zha_gateway.async_block_till_done()
+
+    assert "Failed to remove device" in caplog.text
+    assert "Device removal failed" in caplog.text
+
+
+async def test_gateway_shutdown_group_on_remove_failure(
+    zha_gateway: Gateway,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test that gateway shutdown continues when group.on_remove fails."""
+    zigpy_dev_1 = await zigpy_device_from_json(
+        zha_gateway.application_controller,
+        "tests/data/devices/ikea-of-sweden-tradfri-bulb-gu10-ws-400lm.json",
+    )
+    zha_device_1 = await join_zigpy_device(zha_gateway, zigpy_dev_1)
+
+    zha_group = await zha_gateway.async_create_zigpy_group(
+        "Test Group",
+        [GroupMemberReference(ieee=zha_device_1.ieee, endpoint_id=1)],
+    )
+    await zha_gateway.async_block_till_done()
+
+    with patch.object(
+        zha_group, "on_remove", side_effect=Exception("Group removal failed")
+    ):
+        await zha_gateway.shutdown()
+        await zha_gateway.async_block_till_done()
+
+    assert "Failed to remove group" in caplog.text
+    assert "Group removal failed" in caplog.text
+
+
+async def test_group_on_remove_entity_failure(
+    zha_gateway: Gateway,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test that group.on_remove continues when group entity removal fails."""
+    zigpy_dev_1 = await zigpy_device_from_json(
+        zha_gateway.application_controller,
+        "tests/data/devices/ikea-of-sweden-tradfri-bulb-gu10-ws-400lm.json",
+    )
+    zigpy_dev_2 = await zigpy_device_from_json(
+        zha_gateway.application_controller,
+        "tests/data/devices/ikea-of-sweden-tradfri-bulb-e26-opal-1000lm.json",
+    )
+    zha_device_1 = await join_zigpy_device(zha_gateway, zigpy_dev_1)
+    zha_device_2 = await join_zigpy_device(zha_gateway, zigpy_dev_2)
+
+    members = [
+        GroupMemberReference(ieee=zha_device_1.ieee, endpoint_id=1),
+        GroupMemberReference(ieee=zha_device_2.ieee, endpoint_id=1),
+    ]
+
+    zha_group = await zha_gateway.async_create_zigpy_group("Test Group", members)
+    await zha_gateway.async_block_till_done()
+
+    group_entity = get_group_entity(zha_group, platform=Platform.LIGHT)
+
+    with patch.object(
+        group_entity, "on_remove", side_effect=Exception("Group entity removal failed")
+    ):
+        await zha_group.on_remove()
+
+    assert "Failed to remove group entity" in caplog.text
+    assert "Group entity removal failed" in caplog.text

zha/application/gateway.py

@@ -741,10 +741,24 @@ async def shutdown(self) -> None:
         self._device_availability_checker.stop()
 
         for device in self._devices.values():
-            await device.on_remove()
+            try:
+                await device.on_remove()
+            except Exception:
+                _LOGGER.warning(
+                    "Failed to remove device %s during shutdown",
+                    device,
+                    exc_info=True,
+                )
 
         for group in self._groups.values():
-            await group.on_remove()
+            try:
+                await group.on_remove()
+            except Exception:
+                _LOGGER.warning(
+                    "Failed to remove group %s during shutdown",
+                    group,
+                    exc_info=True,
+                )
 
         _LOGGER.debug("Shutting down ZHA ControllerApplication")
         if self.application_controller is not None:

zha/zigbee/device.py

@@ -1024,13 +1024,37 @@ def entity_update_listener(event: EntityStateChangedEvent) -> None:
     async def on_remove(self) -> None:
         """Cancel tasks this device owns."""
         for callback in self._on_remove_callbacks:
-            callback()
+            try:
+                callback()
+            except Exception:
+                _LOGGER.warning(
+                    "Failed to execute on_remove callback %s for device %s",
+                    callback,
+                    self,
+                    exc_info=True,
+                )
 
         for platform_entity in self._platform_entities.values():
-            await platform_entity.on_remove()
+            try:
+                await platform_entity.on_remove()
+            except Exception:
+                _LOGGER.warning(
+                    "Failed to remove platform entity %s for device %s",
+                    platform_entity,
+                    self,
+                    exc_info=True,
+                )
 
         for entity in self._pending_entities:
-            await entity.on_remove()
+            try:
+                await entity.on_remove()
+            except Exception:
+                _LOGGER.warning(
+                    "Failed to remove pending entity %s for device %s",
+                    entity,
+                    self,
+                    exc_info=True,
+                )
 
     def async_get_clusters(self) -> dict[int, dict[str, dict[int, Cluster]]]:
         """Get all clusters for this device."""

zha/zigbee/group.py

@@ -354,4 +354,11 @@ def log(self, level: int, msg: str, *args: Any, **kwargs) -> None:
     async def on_remove(self) -> None:
         """Cancel tasks this group owns."""
         for group_entity in tuple(self._group_entities.values()):
-            await group_entity.on_remove()
+            try:
+                await group_entity.on_remove()
+            except Exception:
+                _LOGGER.warning(
+                    "Failed to remove group entity %s",
+                    group_entity,
+                    exc_info=True,
+                )