From 3df03f5be513d402a39ef9a64ea081766df324d9 Mon Sep 17 00:00:00 2001
From: Franck Nijhof <git@frenck.dev>
Date: Thu, 28 Mar 2024 16:57:29 +0100
Subject: [PATCH] Fix area search for entities of devices (#114394)

---
 homeassistant/components/search/__init__.py |  9 +++++---
 tests/components/search/test_init.py        | 24 ++++++++++++++++++---
 2 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/homeassistant/components/search/__init__.py b/homeassistant/components/search/__init__.py
index 71b51210a25..a85a21e8102 100644
--- a/homeassistant/components/search/__init__.py
+++ b/homeassistant/components/search/__init__.py
@@ -136,6 +136,9 @@ class Searcher:
         # Scripts referencing this area
         self._add(ItemType.SCRIPT, script.scripts_with_area(self.hass, area_id))
 
+        # Entity in this area, will extend this with the entities of the devices in this area
+        entity_entries = er.async_entries_for_area(self._entity_registry, area_id)
+
         # Devices in this area
         for device in dr.async_entries_for_area(self._device_registry, area_id):
             self._add(ItemType.DEVICE, device.id)
@@ -160,10 +163,10 @@ class Searcher:
                 # Skip the entity if it's in a different area
                 if entity_entry.area_id is not None:
                     continue
-                self._add(ItemType.ENTITY, entity_entry.entity_id)
+                entity_entries.append(entity_entry)
 
-        # Entities in this area
-        for entity_entry in er.async_entries_for_area(self._entity_registry, area_id):
+        # Process entities in this area
+        for entity_entry in entity_entries:
             self._add(ItemType.ENTITY, entity_entry.entity_id)
 
             # If this entity also exists as a resource, we add it.
diff --git a/tests/components/search/test_init.py b/tests/components/search/test_init.py
index ee7b60dc9ac..a817fbfc39e 100644
--- a/tests/components/search/test_init.py
+++ b/tests/components/search/test_init.py
@@ -496,11 +496,14 @@ async def test_search(
         ItemType.SCRIPT: {script_scene_entity.entity_id, "script.nested"},
     }
     assert search(ItemType.AREA, living_room_area.id) == {
-        ItemType.AUTOMATION: {"automation.wled_device"},
+        ItemType.AUTOMATION: {"automation.wled_device", "automation.wled_entity"},
         ItemType.CONFIG_ENTRY: {wled_config_entry.entry_id},
         ItemType.DEVICE: {wled_device.id},
         ItemType.ENTITY: {wled_segment_1_entity.entity_id},
         ItemType.FLOOR: {first_floor.floor_id},
+        ItemType.GROUP: {"group.wled", "group.wled_hue"},
+        ItemType.SCENE: {"scene.scene_wled_seg_1", scene_wled_hue_entity.entity_id},
+        ItemType.SCRIPT: {"script.wled"},
     }
     assert search(ItemType.AREA, kitchen_area.id) == {
         ItemType.AUTOMATION: {"automation.area"},
@@ -511,7 +514,9 @@ async def test_search(
             hue_segment_2_entity.entity_id,
         },
         ItemType.FLOOR: {first_floor.floor_id},
-        ItemType.SCRIPT: {"script.area", "script.device"},
+        ItemType.GROUP: {"group.hue", "group.wled_hue"},
+        ItemType.SCENE: {"scene.scene_hue_seg_1", scene_wled_hue_entity.entity_id},
+        ItemType.SCRIPT: {"script.area", "script.device", "script.hue"},
     }
 
     assert not search(ItemType.AUTOMATION, "automation.unknown")
@@ -726,6 +731,7 @@ async def test_search(
             "automation.area",
             "automation.floor",
             "automation.wled_device",
+            "automation.wled_entity",
         },
         ItemType.CONFIG_ENTRY: {hue_config_entry.entry_id, wled_config_entry.entry_id},
         ItemType.DEVICE: {hue_device.id, wled_device.id},
@@ -734,7 +740,19 @@ async def test_search(
             hue_segment_1_entity.entity_id,
             hue_segment_2_entity.entity_id,
         },
-        ItemType.SCRIPT: {"script.device", "script.area", "script.floor"},
+        ItemType.GROUP: {"group.hue", "group.wled", "group.wled_hue"},
+        ItemType.SCENE: {
+            "scene.scene_hue_seg_1",
+            "scene.scene_wled_seg_1",
+            scene_wled_hue_entity.entity_id,
+        },
+        ItemType.SCRIPT: {
+            "script.device",
+            "script.area",
+            "script.floor",
+            "script.hue",
+            "script.wled",
+        },
     }
     assert search(ItemType.FLOOR, second_floor.floor_id) == {
         ItemType.AREA: {bedroom_area.id},