From 69f54656c44b9de30777abdce925cd0bb88be76a Mon Sep 17 00:00:00 2001
From: Michael <35783820+mib1185@users.noreply.github.com>
Date: Wed, 31 Jul 2024 17:58:11 +0200
Subject: [PATCH] Fix cleanup of orphan device entries in AVM Fritz!Box Tools
 (#122937)

* fix cleanup of orphan device entries

* add test for cleanup button
---
 homeassistant/components/fritz/coordinator.py |  4 +-
 tests/components/fritz/test_button.py         | 74 +++++++++++++++++--
 2 files changed, 72 insertions(+), 6 deletions(-)

diff --git a/homeassistant/components/fritz/coordinator.py b/homeassistant/components/fritz/coordinator.py
index 8a55084d7ef..592bf37084e 100644
--- a/homeassistant/components/fritz/coordinator.py
+++ b/homeassistant/components/fritz/coordinator.py
@@ -666,7 +666,9 @@ class FritzBoxTools(DataUpdateCoordinator[UpdateCoordinatorDataType]):
                 entity_reg.async_remove(entity.entity_id)
 
         device_reg = dr.async_get(self.hass)
-        orphan_connections = {(CONNECTION_NETWORK_MAC, mac) for mac in orphan_macs}
+        orphan_connections = {
+            (CONNECTION_NETWORK_MAC, dr.format_mac(mac)) for mac in orphan_macs
+        }
         for device in dr.async_entries_for_config_entry(
             device_reg, config_entry.entry_id
         ):
diff --git a/tests/components/fritz/test_button.py b/tests/components/fritz/test_button.py
index 8666491eb7a..79639835003 100644
--- a/tests/components/fritz/test_button.py
+++ b/tests/components/fritz/test_button.py
@@ -1,6 +1,6 @@
 """Tests for Fritz!Tools button platform."""
 
-import copy
+from copy import deepcopy
 from datetime import timedelta
 from unittest.mock import patch
 
@@ -11,9 +11,15 @@ from homeassistant.components.fritz.const import DOMAIN, MeshRoles
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.const import ATTR_ENTITY_ID, STATE_UNKNOWN
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.util.dt import utcnow
 
-from .const import MOCK_MESH_DATA, MOCK_NEW_DEVICE_NODE, MOCK_USER_DATA
+from .const import (
+    MOCK_HOST_ATTRIBUTES_DATA,
+    MOCK_MESH_DATA,
+    MOCK_NEW_DEVICE_NODE,
+    MOCK_USER_DATA,
+)
 
 from tests.common import MockConfigEntry, async_fire_time_changed
 
@@ -120,7 +126,7 @@ async def test_wol_button_new_device(
     entry = MockConfigEntry(domain=DOMAIN, data=MOCK_USER_DATA)
     entry.add_to_hass(hass)
 
-    mesh_data = copy.deepcopy(MOCK_MESH_DATA)
+    mesh_data = deepcopy(MOCK_MESH_DATA)
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
     assert entry.state is ConfigEntryState.LOADED
@@ -148,7 +154,7 @@ async def test_wol_button_absent_for_mesh_slave(
     entry = MockConfigEntry(domain=DOMAIN, data=MOCK_USER_DATA)
     entry.add_to_hass(hass)
 
-    slave_mesh_data = copy.deepcopy(MOCK_MESH_DATA)
+    slave_mesh_data = deepcopy(MOCK_MESH_DATA)
     slave_mesh_data["nodes"][0]["mesh_role"] = MeshRoles.SLAVE
     fh_class_mock.get_mesh_topology.return_value = slave_mesh_data
 
@@ -170,7 +176,7 @@ async def test_wol_button_absent_for_non_lan_device(
     entry = MockConfigEntry(domain=DOMAIN, data=MOCK_USER_DATA)
     entry.add_to_hass(hass)
 
-    printer_wifi_data = copy.deepcopy(MOCK_MESH_DATA)
+    printer_wifi_data = deepcopy(MOCK_MESH_DATA)
     # initialization logic uses the connection type of the `node_interface_1_uid` pair of the printer
     # ni-230 is wifi interface of fritzbox
     printer_node_interface = printer_wifi_data["nodes"][1]["node_interfaces"][0]
@@ -184,3 +190,61 @@ async def test_wol_button_absent_for_non_lan_device(
 
     button = hass.states.get("button.printer_wake_on_lan")
     assert button is None
+
+
+async def test_cleanup_button(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
+    fc_class_mock,
+    fh_class_mock,
+) -> None:
+    """Test cleanup of orphan devices."""
+
+    entry = MockConfigEntry(domain=DOMAIN, data=MOCK_USER_DATA)
+    entry.add_to_hass(hass)
+
+    await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
+    assert entry.state is ConfigEntryState.LOADED
+
+    # check if tracked device is registered properly
+    device = device_registry.async_get_device(
+        connections={("mac", "aa:bb:cc:00:11:22")}
+    )
+    assert device
+
+    entities = [
+        entity
+        for entity in er.async_entries_for_config_entry(entity_registry, entry.entry_id)
+        if entity.unique_id.startswith("AA:BB:CC:00:11:22")
+    ]
+    assert entities
+    assert len(entities) == 3
+
+    # removed tracked device and trigger cleanup
+    host_attributes = deepcopy(MOCK_HOST_ATTRIBUTES_DATA)
+    host_attributes.pop(0)
+    fh_class_mock.get_hosts_attributes.return_value = host_attributes
+
+    await hass.services.async_call(
+        BUTTON_DOMAIN,
+        SERVICE_PRESS,
+        {ATTR_ENTITY_ID: "button.mock_title_cleanup"},
+        blocking=True,
+    )
+
+    await hass.async_block_till_done(wait_background_tasks=True)
+
+    # check if orphan tracked device is removed
+    device = device_registry.async_get_device(
+        connections={("mac", "aa:bb:cc:00:11:22")}
+    )
+    assert not device
+
+    entities = [
+        entity
+        for entity in er.async_entries_for_config_entry(entity_registry, entry.entry_id)
+        if entity.unique_id.startswith("AA:BB:CC:00:11:22")
+    ]
+    assert not entities