Fix shelly remove orphaned entities (home-assistant#154182)

Co-authored-by: Franck Nijhof <[email protected]>

Author: thecode

homeassistant/components/shelly/utils.py

@@ -692,27 +692,25 @@ def async_remove_orphaned_entities(
     """Remove orphaned entities."""
     orphaned_entities = []
     entity_reg = er.async_get(hass)
-    device_reg = dr.async_get(hass)
 
-    if not (
-        devices := device_reg.devices.get_devices_for_config_entry_id(config_entry_id)
-    ):
-        return
+    entities = er.async_entries_for_config_entry(entity_reg, config_entry_id)
+    for entity in entities:
+        if not entity.entity_id.startswith(platform):
+            continue
+        if key_suffix is not None and key_suffix not in entity.unique_id:
+            continue
+        # we are looking for the component ID, e.g. boolean:201, em1data:1
+        if not (match := COMPONENT_ID_PATTERN.search(entity.unique_id)):
+            continue
 
-    for device in devices:
-        entities = er.async_entries_for_device(entity_reg, device.id, True)
-        for entity in entities:
-            if not entity.entity_id.startswith(platform):
-                continue
-            if key_suffix is not None and key_suffix not in entity.unique_id:
-                continue
-            # we are looking for the component ID, e.g. boolean:201, em1data:1
-            if not (match := COMPONENT_ID_PATTERN.search(entity.unique_id)):
-                continue
-
-            key = match.group()
-            if key not in keys:
-                orphaned_entities.append(entity.unique_id.split("-", 1)[1])
+        key = match.group()
+        if key not in keys:
+            LOGGER.debug(
+                "Found orphaned Shelly entity: %s, unique id: %s",
+                entity.entity_id,
+                entity.unique_id,
+            )
+            orphaned_entities.append(entity.unique_id.split("-", 1)[1])
 
     if orphaned_entities:
         async_remove_shelly_rpc_entities(hass, platform, mac, orphaned_entities)

tests/components/shelly/test_sensor.py

@@ -1088,6 +1088,7 @@ async def test_rpc_device_virtual_text_sensor(
         ("text", "text_generic", None),
     ],
 )
+@pytest.mark.usefixtures("disable_async_remove_shelly_rpc_entities")
 async def test_migrate_unique_id_virtual_components_roles(
     hass: HomeAssistant,
     mock_rpc_device: Mock,