Use _remove_entity to remove entities

Author: puddly

tests/test_discover.py

@@ -844,8 +844,6 @@ async def test_devices_from_files(
 
                 unique_id_migrations[key] = entity
 
-        await zha_device.on_remove()
-
         # XXX: We re-serialize the JSON because integer enum types are converted when
         # serializing but will not compare properly otherwise
         loaded_device_data = json.loads(
@@ -874,3 +872,5 @@ async def test_devices_from_files(
                     tsn=None,
                 )
             ]
+
+        await zha_device.on_remove()

zha/zigbee/device.py

@@ -1111,8 +1111,10 @@ async def on_remove(self) -> None:
         for callback in self._on_remove_callbacks:
             callback()
 
-        for platform_entity in self._platform_entities.values():
-            await platform_entity.on_remove()
+        for platform_entity in list(self._platform_entities.values()):
+            # TODO: To avoid unnecessary traffic during shutdown, we don't need to emit
+            # an event for every entity, just the device
+            await self._remove_entity(platform_entity, emit_event=False)
 
         for entity in self._pending_entities:
             await entity.on_remove()