fix autoremove

dmamontov · Mar 13, 2021 · 05cf796 · 05cf796
1 parent f45b0e0
commit 05cf796
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 3 deletions.
diff --git a/README.md b/README.md
@@ -109,7 +109,7 @@ target:
 ## Routers tested
 Many more Xiaomi and Redmi routers supported by MiWiFi (OpenWRT - Luci API)
 
-| Image                                               | Router                                                                              | Firmware version           | Status                        |
+| Image                                               | Router                                                           | Firmware version           | Status                        |
 | --------------------------------------------------- | ---------------------------------------------------------------- | -------------------------- | ----------------------------- |
 | ![](http://www1.miwifi.com/statics/img/2100@1x.png) | [Xiaomi AC2100](https://www.mi.com/miwifiac)                     | 2.0.743(CN)                | Supported                     |
 | ![](http://www1.miwifi.com/statics/img/RA72.png)    | [Xiaomi AX3600](https://www.mi.com/r3600)                        | 1.0.79(CN), 3.0.22(Global) | Supported                     |

diff --git a/custom_components/miwifi/__init__.py b/custom_components/miwifi/__init__.py
@@ -86,9 +86,9 @@ async def remove_devices(call: ServiceCall):
         devices = data.pop('device_id', [])
         entities = data.pop('entity_id', None)
 
-        if entities:
-            entity_registry = await er.async_get_registry(hass)
+        entity_registry = await er.async_get_registry(hass)
 
+        if entities:
             for entity_id in entities:
                 entity = entity_registry.async_get(entity_id)
                 if not entity:
@@ -104,6 +104,11 @@ async def remove_devices(call: ServiceCall):
                 if not device_entry:
                     continue
 
+                device_entities = er.async_entries_for_device(entity_registry, device, True)
+                if device_entities:
+                    for entity in device_entities:
+                        entity_registry.async_remove(entity.entity_id)
+
                 device_registry.async_remove_device(device)
 
                 for entry in device_entry.config_entries:

diff --git a/custom_components/miwifi/core/luci_data.py b/custom_components/miwifi/core/luci_data.py
@@ -2,6 +2,7 @@
 
 import asyncio
 import homeassistant.helpers.device_registry as dr
+import homeassistant.helpers.entity_registry as er
 
 from datetime import datetime, timedelta
 
@@ -162,6 +163,7 @@ async def check_last_activity(self) -> None:
             return
 
         device_registry = await dr.async_get_registry(self.hass)
+        entity_registry = await er.async_get_registry(self.hass)
         now = datetime.now().replace(microsecond = 0)
 
         for mac in self._devices:
@@ -179,6 +181,11 @@ async def check_last_activity(self) -> None:
                 if not device:
                     continue
 
+                device_entities = er.async_entries_for_device(entity_registry, device.id, True)
+                if device_entities:
+                    for entity in device_entities:
+                        entity_registry.async_remove(entity.entity_id)
+
                 device_registry.async_remove_device(device.id)
 
     async def save_to_store(self) -> None: