From c3a432620c0b537f176fe02be2880ba85a77e905 Mon Sep 17 00:00:00 2001
From: Simone Chemelli <simone.chemelli@gmail.com>
Date: Wed, 5 Jan 2022 20:21:15 +0100
Subject: [PATCH] Improve cleanup service for Fritz (#61484)

---
 homeassistant/components/fritz/common.py | 30 ++++++++++++++++++------
 homeassistant/components/fritz/switch.py | 11 +++++++--
 2 files changed, 32 insertions(+), 9 deletions(-)

diff --git a/homeassistant/components/fritz/common.py b/homeassistant/components/fritz/common.py
index e62c56f6f2d..e580665cfc5 100644
--- a/homeassistant/components/fritz/common.py
+++ b/homeassistant/components/fritz/common.py
@@ -426,8 +426,9 @@ class FritzBoxTools(update_coordinator.DataUpdateCoordinator):
                 )
                 return
 
+            device_hosts_list: list[dict] = []
             if service_call.service == SERVICE_CLEANUP:
-                device_hosts_list: list = await self.hass.async_add_executor_job(
+                device_hosts_list = await self.hass.async_add_executor_job(
                     self.fritz_hosts.get_hosts_info
                 )
 
@@ -447,15 +448,30 @@ class FritzBoxTools(update_coordinator.DataUpdateCoordinator):
         )
         entities_removed: bool = False
 
-        device_hosts_macs = {device["mac"] for device in device_hosts_list}
+        device_hosts_macs = set()
+        device_hosts_names = set()
+        for device in device_hosts_list:
+            device_hosts_macs.add(device["mac"])
+            device_hosts_names.add(device["name"])
 
         for entry in ha_entity_reg_list:
-            if (
-                not _cleanup_entity_filter(entry)
-                or entry.unique_id.split("_")[0] in device_hosts_macs
-            ):
+            if entry.original_name is None:
                 continue
-            _LOGGER.info("Removing entity: %s", entry.name or entry.original_name)
+            entry_name = entry.name or entry.original_name
+            entry_host = entry_name.split(" ")[0]
+            entry_mac = entry.unique_id.split("_")[0]
+
+            if not _cleanup_entity_filter(entry) or (
+                entry_mac in device_hosts_macs and entry_host in device_hosts_names
+            ):
+                _LOGGER.debug(
+                    "Skipping entity %s [mac=%s, host=%s]",
+                    entry_name,
+                    entry_mac,
+                    entry_host,
+                )
+                continue
+            _LOGGER.info("Removing entity: %s", entry_name)
             entity_reg.async_remove(entry.entity_id)
             entities_removed = True
 
diff --git a/homeassistant/components/fritz/switch.py b/homeassistant/components/fritz/switch.py
index 773b9c771b7..1946ed7fb19 100644
--- a/homeassistant/components/fritz/switch.py
+++ b/homeassistant/components/fritz/switch.py
@@ -10,6 +10,7 @@ from fritzconnection.core.exceptions import (
     FritzActionError,
     FritzActionFailedError,
     FritzConnectionException,
+    FritzLookUpError,
     FritzSecurityError,
     FritzServiceError,
 )
@@ -90,10 +91,16 @@ def service_call_action(
             exc_info=True,
         )
         return None
-    except (FritzActionError, FritzActionFailedError, FritzServiceError):
+    except (
+        FritzActionError,
+        FritzActionFailedError,
+        FritzServiceError,
+        FritzLookUpError,
+    ):
         _LOGGER.error(
-            "Service/Action Error: cannot execute service %s",
+            "Service/Action Error: cannot execute service %s with action %s",
             service_name,
+            action_name,
             exc_info=True,
         )
         return None