Fix Shelly orphaned entity removal logic to handle sub-devices (#152195)

homeassistant/components/shelly/utils.py

@@ -682,20 +682,20 @@ def async_remove_orphaned_entities(
     ):
         return
 
-    device_id = devices[0].id
-    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
+    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:
+                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/__init__.py

@@ -156,6 +156,17 @@ def register_device(
     )
 
 
+def register_sub_device(
+    device_registry: DeviceRegistry, config_entry: ConfigEntry, unique_id: str
+) -> DeviceEntry:
+    """Register Shelly sub-device."""
+    return device_registry.async_get_or_create(
+        config_entry_id=config_entry.entry_id,
+        identifiers={(DOMAIN, f"{MOCK_MAC}-{unique_id}")},
+        via_device=(DOMAIN, format_mac(MOCK_MAC)),
+    )
+
+
 async def snapshot_device_entities(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,

tests/components/shelly/test_binary_sensor.py

@@ -21,6 +21,7 @@ from . import (
     mutate_rpc_device_status,
     register_device,
     register_entity,
+    register_sub_device,
 )
 
 from tests.common import mock_restore_cache
@@ -475,8 +476,10 @@ async def test_rpc_remove_virtual_binary_sensor_when_orphaned(
 ) -> None:
     """Check whether the virtual binary sensor will be removed if it has been removed from the device configuration."""
     config_entry = await init_integration(hass, 3, skip_setup=True)
+
+    # create orphaned entity on main device
     device_entry = register_device(device_registry, config_entry)
-    entity_id = register_entity(
+    entity_id1 = register_entity(
         hass,
         BINARY_SENSOR_DOMAIN,
         "test_name_boolean_200",
@@ -485,10 +488,29 @@ async def test_rpc_remove_virtual_binary_sensor_when_orphaned(
         device_id=device_entry.id,
     )
 
+    # create orphaned entity on sub device
+    sub_device_entry = register_sub_device(
+        device_registry,
+        config_entry,
+        "boolean:201-boolean",
+    )
+    entity_id2 = register_entity(
+        hass,
+        BINARY_SENSOR_DOMAIN,
+        "boolean_201",
+        "boolean:201-boolean",
+        config_entry,
+        device_id=sub_device_entry.id,
+    )
+
+    assert entity_registry.async_get(entity_id1) is not None
+    assert entity_registry.async_get(entity_id2) is not None
+
     await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()
 
-    assert entity_registry.async_get(entity_id) is None
+    assert entity_registry.async_get(entity_id1) is None
+    assert entity_registry.async_get(entity_id2) is None
 
 
 async def test_blu_trv_binary_sensor_entity(

tests/components/shelly/test_switch.py

@@ -36,6 +36,7 @@ from . import (
     inject_rpc_device_event,
     register_device,
     register_entity,
+    register_sub_device,
 )
 
 from tests.common import async_fire_time_changed, mock_restore_cache
@@ -720,8 +721,10 @@ async def test_rpc_remove_virtual_switch_when_orphaned(
 ) -> None:
     """Check whether the virtual switch will be removed if it has been removed from the device configuration."""
     config_entry = await init_integration(hass, 3, skip_setup=True)
+
+    # create orphaned entity on main device
     device_entry = register_device(device_registry, config_entry)
-    entity_id = register_entity(
+    entity_id1 = register_entity(
         hass,
         SWITCH_DOMAIN,
         "test_name_boolean_200",
@@ -730,10 +733,29 @@ async def test_rpc_remove_virtual_switch_when_orphaned(
         device_id=device_entry.id,
     )
 
+    # create orphaned entity on sub device
+    sub_device_entry = register_sub_device(
+        device_registry,
+        config_entry,
+        "boolean:201-boolean",
+    )
+    entity_id2 = register_entity(
+        hass,
+        SWITCH_DOMAIN,
+        "boolean_201",
+        "boolean:201-boolean",
+        config_entry,
+        device_id=sub_device_entry.id,
+    )
+
+    assert entity_registry.async_get(entity_id1) is not None
+    assert entity_registry.async_get(entity_id2) is not None
+
     await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()
 
-    assert entity_registry.async_get(entity_id) is None
+    assert entity_registry.async_get(entity_id1) is None
+    assert entity_registry.async_get(entity_id2) is None
 
 
 @pytest.mark.usefixtures("entity_registry_enabled_by_default")