Add search integration (#30511)

* Add search integration

* Add scenes and config entry support

* Update comments

* Add support for groups

* Allow querying config entry

* Update manifest

* Fix scene tests

CODEOWNERS

@@ -281,6 +281,7 @@ homeassistant/components/samsungtv/* @escoand
 homeassistant/components/scene/* @home-assistant/core
 homeassistant/components/scrape/* @fabaff
 homeassistant/components/script/* @home-assistant/core
+homeassistant/components/search/* @home-assistant/core
 homeassistant/components/sense/* @kbickar
 homeassistant/components/sensibo/* @andrey-git
 homeassistant/components/sentry/* @dcramer

homeassistant/components/group/__init__.py

@@ -183,6 +183,24 @@ def get_entity_ids(
     return [ent_id for ent_id in entity_ids if ent_id.startswith(domain_filter)]
 
 
+@bind_hass
+def groups_with_entity(hass: HomeAssistantType, entity_id: str) -> List[str]:
+    """Get all groups that contain this entity.
+
+    Async friendly.
+    """
+    if DOMAIN not in hass.data:
+        return []
+
+    groups = []
+
+    for group in hass.data[DOMAIN].entities:
+        if entity_id in group.tracking:
+            groups.append(group.entity_id)
+
+    return groups
+
+
 async def async_setup(hass, config):
     """Set up all groups found defined in the configuration."""
     component = hass.data.get(DOMAIN)

homeassistant/components/homeassistant/scene.py

@@ -1,6 +1,7 @@
 """Allow users to set and activate scenes."""
 from collections import namedtuple
 import logging
+from typing import List
 
 import voluptuous as vol
 
@@ -17,7 +18,7 @@ from homeassistant.const import (
     STATE_OFF,
     STATE_ON,
 )
-from homeassistant.core import DOMAIN as HA_DOMAIN, State
+from homeassistant.core import DOMAIN as HA_DOMAIN, HomeAssistant, State, callback
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import (
     config_per_platform,
@@ -71,7 +72,7 @@ def _ensure_no_intersection(value):
 
 CONF_SCENE_ID = "scene_id"
 CONF_SNAPSHOT = "snapshot_entities"
-
+DATA_PLATFORM = f"homeassistant_scene"
 STATES_SCHEMA = vol.All(dict, _convert_states)
 
 PLATFORM_SCHEMA = vol.Schema(
@@ -108,6 +109,39 @@ SCENECONFIG = namedtuple("SceneConfig", [CONF_NAME, STATES])
 _LOGGER = logging.getLogger(__name__)
 
 
+@callback
+def scenes_with_entity(hass: HomeAssistant, entity_id: str) -> List[str]:
+    """Return all scenes that reference the entity."""
+    if DATA_PLATFORM not in hass.data:
+        return []
+
+    platform = hass.data[DATA_PLATFORM]
+
+    results = []
+
+    for scene_entity in platform.entities.values():
+        if entity_id in scene_entity.scene_config.states:
+            results.append(scene_entity.entity_id)
+
+    return results
+
+
+@callback
+def entities_in_scene(hass: HomeAssistant, entity_id: str) -> List[str]:
+    """Return all entities in a scene."""
+    if DATA_PLATFORM not in hass.data:
+        return []
+
+    platform = hass.data[DATA_PLATFORM]
+
+    entity = platform.entities.get(entity_id)
+
+    if entity is None:
+        return []
+
+    return list(entity.scene_config.states)
+
+
 async def async_setup_platform(hass, config, async_add_entities, discovery_info=None):
     """Set up Home Assistant scene entries."""
     _process_scenes_config(hass, async_add_entities, config)
@@ -117,7 +151,7 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
         return
 
     # Store platform for later.
-    platform = entity_platform.current_platform.get()
+    platform = hass.data[DATA_PLATFORM] = entity_platform.current_platform.get()
 
     async def reload_config(call):
         """Reload the scene config."""

homeassistant/components/search/__init__.py

@@ -0,0 +1,211 @@
+"""The Search integration."""
+from collections import defaultdict
+
+import voluptuous as vol
+
+from homeassistant.components import group, websocket_api
+from homeassistant.components.homeassistant import scene
+from homeassistant.core import HomeAssistant, callback, split_entity_id
+from homeassistant.helpers import device_registry, entity_registry
+
+DOMAIN = "search"
+
+
+async def async_setup(hass: HomeAssistant, config: dict):
+    """Set up the Search component."""
+    websocket_api.async_register_command(hass, websocket_search_related)
+    return True
+
+
+@websocket_api.async_response
+@websocket_api.websocket_command(
+    {
+        vol.Required("type"): "search/related",
+        vol.Required("item_type"): vol.In(
+            (
+                "area",
+                "automation",
+                "config_entry",
+                "device",
+                "entity",
+                "group",
+                "scene",
+                "script",
+            )
+        ),
+        vol.Required("item_id"): str,
+    }
+)
+async def websocket_search_related(hass, connection, msg):
+    """Handle search."""
+    searcher = Searcher(
+        hass,
+        await device_registry.async_get_registry(hass),
+        await entity_registry.async_get_registry(hass),
+    )
+    connection.send_result(
+        msg["id"], searcher.async_search(msg["item_type"], msg["item_id"])
+    )
+
+
+class Searcher:
+    """Find related things.
+
+    Few rules:
+    Scenes, scripts, automations and config entries will only be expanded if they are
+    the entry point. They won't be expanded if we process them. This is because they
+    turn the results into garbage.
+    """
+
+    # These types won't be further explored. Config entries + Output types.
+    DONT_RESOLVE = {"scene", "automation", "script", "group", "config_entry"}
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        device_reg: device_registry.DeviceRegistry,
+        entity_reg: entity_registry.EntityRegistry,
+    ):
+        """Search results."""
+        self.hass = hass
+        self._device_reg = device_reg
+        self._entity_reg = entity_reg
+        self.results = defaultdict(set)
+        self._to_resolve = set()
+
+    @callback
+    def async_search(self, item_type, item_id):
+        """Find results."""
+        self.results[item_type].add(item_id)
+        self._to_resolve.add((item_type, item_id))
+
+        while self._to_resolve:
+            search_type, search_id = self._to_resolve.pop()
+            getattr(self, f"_resolve_{search_type}")(search_id)
+
+        # Clean up entity_id items, from the general "entity" type result,
+        # that are also found in the specific entity domain type.
+        self.results["entity"] -= self.results["script"]
+        self.results["entity"] -= self.results["scene"]
+        self.results["entity"] -= self.results["automation"]
+        self.results["entity"] -= self.results["group"]
+
+        # Remove entry into graph from search results.
+        self.results[item_type].remove(item_id)
+
+        # Filter out empty sets.
+        return {key: val for key, val in self.results.items() if val}
+
+    @callback
+    def _add_or_resolve(self, item_type, item_id):
+        """Add an item to explore."""
+        if item_id in self.results[item_type]:
+            return
+
+        self.results[item_type].add(item_id)
+
+        if item_type not in self.DONT_RESOLVE:
+            self._to_resolve.add((item_type, item_id))
+
+    @callback
+    def _resolve_area(self, area_id) -> None:
+        """Resolve an area."""
+        for device in device_registry.async_entries_for_area(self._device_reg, area_id):
+            self._add_or_resolve("device", device.id)
+
+    @callback
+    def _resolve_device(self, device_id) -> None:
+        """Resolve a device."""
+        device_entry = self._device_reg.async_get(device_id)
+        # Unlikely entry doesn't exist, but let's guard for bad data.
+        if device_entry is not None:
+            if device_entry.area_id:
+                self._add_or_resolve("area", device_entry.area_id)
+
+            for config_entry_id in device_entry.config_entries:
+                self._add_or_resolve("config_entry", config_entry_id)
+
+            # We do not resolve device_entry.via_device_id because that
+            # device is not related data-wise inside HA.
+
+        for entity_entry in entity_registry.async_entries_for_device(
+            self._entity_reg, device_id
+        ):
+            self._add_or_resolve("entity", entity_entry.entity_id)
+
+        # Extra: Find automations that reference this device
+
+    @callback
+    def _resolve_entity(self, entity_id) -> None:
+        """Resolve an entity."""
+        # Extra: Find automations and scripts that reference this entity.
+
+        for entity in scene.scenes_with_entity(self.hass, entity_id):
+            self._add_or_resolve("entity", entity)
+
+        for entity in group.groups_with_entity(self.hass, entity_id):
+            self._add_or_resolve("entity", entity)
+
+        # Find devices
+        entity_entry = self._entity_reg.async_get(entity_id)
+        if entity_entry is not None:
+            if entity_entry.device_id:
+                self._add_or_resolve("device", entity_entry.device_id)
+
+            if entity_entry.config_entry_id is not None:
+                self._add_or_resolve("config_entry", entity_entry.config_entry_id)
+
+        domain = split_entity_id(entity_id)[0]
+
+        if domain in ("scene", "automation", "script", "group"):
+            self._add_or_resolve(domain, entity_id)
+
+    @callback
+    def _resolve_automation(self, automation_entity_id) -> None:
+        """Resolve an automation.
+
+        Will only be called if automation is an entry point.
+        """
+        # Extra: Check with automation integration what entities/devices they reference
+
+    @callback
+    def _resolve_script(self, script_entity_id) -> None:
+        """Resolve a script.
+
+        Will only be called if script is an entry point.
+        """
+        # Extra: Check with script integration what entities/devices they reference
+
+    @callback
+    def _resolve_group(self, group_entity_id) -> None:
+        """Resolve a group.
+
+        Will only be called if group is an entry point.
+        """
+        for entity_id in group.get_entity_ids(self.hass, group_entity_id):
+            self._add_or_resolve("entity", entity_id)
+
+    @callback
+    def _resolve_scene(self, scene_entity_id) -> None:
+        """Resolve a scene.
+
+        Will only be called if scene is an entry point.
+        """
+        for entity in scene.entities_in_scene(self.hass, scene_entity_id):
+            self._add_or_resolve("entity", entity)
+
+    @callback
+    def _resolve_config_entry(self, config_entry_id) -> None:
+        """Resolve a config entry.
+
+        Will only be called if config entry is an entry point.
+        """
+        for device_entry in device_registry.async_entries_for_config_entry(
+            self._device_reg, config_entry_id
+        ):
+            self._add_or_resolve("device", device_entry.id)
+
+        for entity_entry in entity_registry.async_entries_for_config_entry(
+            self._entity_reg, config_entry_id
+        ):
+            self._add_or_resolve("entity", entity_entry.entity_id)

homeassistant/components/search/manifest.json

@@ -0,0 +1,12 @@
+{
+  "domain": "search",
+  "name": "Search",
+  "documentation": "https://www.home-assistant.io/integrations/search",
+  "requirements": [],
+  "ssdp": [],
+  "zeroconf": [],
+  "homekit": {},
+  "dependencies": ["websocket_api"],
+  "after_dependencies": ["scene", "group"],
+  "codeowners": ["@home-assistant/core"]
+}

homeassistant/helpers/device_registry.py

@@ -375,3 +375,15 @@ async def async_get_registry(hass: HomeAssistantType) -> DeviceRegistry:
 def async_entries_for_area(registry: DeviceRegistry, area_id: str) -> List[DeviceEntry]:
     """Return entries that match an area."""
     return [device for device in registry.devices.values() if device.area_id == area_id]
+
+
+@callback
+def async_entries_for_config_entry(
+    registry: DeviceRegistry, config_entry_id: str
+) -> List[DeviceEntry]:
+    """Return entries that match a config entry."""
+    return [
+        device
+        for device in registry.devices.values()
+        if config_entry_id in device.config_entries
+    ]

homeassistant/helpers/entity_registry.py

@@ -445,6 +445,18 @@ def async_entries_for_device(
     ]
 
 
+@callback
+def async_entries_for_config_entry(
+    registry: EntityRegistry, config_entry_id: str
+) -> List[RegistryEntry]:
+    """Return entries that match a config entry."""
+    return [
+        entry
+        for entry in registry.entities.values()
+        if entry.config_entry_id == config_entry_id
+    ]
+
+
 async def _async_migrate(entities: Dict[str, Any]) -> Dict[str, List[Dict[str, Any]]]:
     """Migrate the YAML config file to storage helper format."""
     return {

tests/components/homeassistant/test_scene.py

@@ -4,6 +4,7 @@ from unittest.mock import patch
 import pytest
 import voluptuous as vol
 
+from homeassistant.components.homeassistant import scene as ha_scene
 from homeassistant.setup import async_setup_component
 
 from tests.common import async_mock_service
@@ -209,3 +210,51 @@ async def test_ensure_no_intersection(hass):
         await hass.async_block_till_done()
     assert "entities and snapshot_entities must not overlap" in str(ex.value)
     assert hass.states.get("scene.hallo") is None
+
+
+async def test_scenes_with_entity(hass):
+    """Test finding scenes with a specific entity."""
+    assert await async_setup_component(
+        hass,
+        "scene",
+        {
+            "scene": [
+                {"name": "scene_1", "entities": {"light.kitchen": "on"}},
+                {"name": "scene_2", "entities": {"light.living_room": "off"}},
+                {
+                    "name": "scene_3",
+                    "entities": {"light.kitchen": "on", "light.living_room": "off"},
+                },
+            ]
+        },
+    )
+
+    assert ha_scene.scenes_with_entity(hass, "light.kitchen") == [
+        "scene.scene_1",
+        "scene.scene_3",
+    ]
+
+
+async def test_entities_in_scene(hass):
+    """Test finding entities in a scene."""
+    assert await async_setup_component(
+        hass,
+        "scene",
+        {
+            "scene": [
+                {"name": "scene_1", "entities": {"light.kitchen": "on"}},
+                {"name": "scene_2", "entities": {"light.living_room": "off"}},
+                {
+                    "name": "scene_3",
+                    "entities": {"light.kitchen": "on", "light.living_room": "off"},
+                },
+            ]
+        },
+    )
+
+    for scene_id, entities in (
+        ("scene.scene_1", ["light.kitchen"]),
+        ("scene.scene_2", ["light.living_room"]),
+        ("scene.scene_3", ["light.kitchen", "light.living_room"]),
+    ):
+        assert ha_scene.entities_in_scene(hass, scene_id) == entities

tests/components/search/__init__.py

@@ -0,0 +1 @@
+"""Tests for the Search integration."""

tests/components/search/test_init.py

@@ -0,0 +1,228 @@
+"""Tests for Search integration."""
+from homeassistant.components import search
+from homeassistant.setup import async_setup_component
+
+from tests.common import MockConfigEntry
+
+
+async def test_search(hass):
+    """Test that search works."""
+    area_reg = await hass.helpers.area_registry.async_get_registry()
+    device_reg = await hass.helpers.device_registry.async_get_registry()
+    entity_reg = await hass.helpers.entity_registry.async_get_registry()
+
+    living_room_area = area_reg.async_create("Living Room")
+
+    # Light strip with 2 lights.
+    wled_config_entry = MockConfigEntry(domain="wled")
+    wled_config_entry.add_to_hass(hass)
+
+    wled_device = device_reg.async_get_or_create(
+        config_entry_id=wled_config_entry.entry_id,
+        name="Light Strip",
+        identifiers=({"wled", "wled-1"}),
+    )
+
+    device_reg.async_update_device(wled_device.id, area_id=living_room_area.id)
+
+    wled_segment_1_entity = entity_reg.async_get_or_create(
+        "light",
+        "wled",
+        "wled-1-seg-1",
+        suggested_object_id="wled segment 1",
+        config_entry=wled_config_entry,
+        device_id=wled_device.id,
+    )
+    wled_segment_2_entity = entity_reg.async_get_or_create(
+        "light",
+        "wled",
+        "wled-1-seg-2",
+        suggested_object_id="wled segment 2",
+        config_entry=wled_config_entry,
+        device_id=wled_device.id,
+    )
+
+    # Non related info.
+    kitchen_area = area_reg.async_create("Kitchen")
+
+    hue_config_entry = MockConfigEntry(domain="hue")
+    hue_config_entry.add_to_hass(hass)
+
+    hue_device = device_reg.async_get_or_create(
+        config_entry_id=hue_config_entry.entry_id,
+        name="Light Strip",
+        identifiers=({"hue", "hue-1"}),
+    )
+
+    device_reg.async_update_device(hue_device.id, area_id=kitchen_area.id)
+
+    hue_segment_1_entity = entity_reg.async_get_or_create(
+        "light",
+        "hue",
+        "hue-1-seg-1",
+        suggested_object_id="hue segment 1",
+        config_entry=hue_config_entry,
+        device_id=hue_device.id,
+    )
+    hue_segment_2_entity = entity_reg.async_get_or_create(
+        "light",
+        "hue",
+        "hue-1-seg-2",
+        suggested_object_id="hue segment 2",
+        config_entry=hue_config_entry,
+        device_id=hue_device.id,
+    )
+
+    await async_setup_component(
+        hass,
+        "group",
+        {
+            "group": {
+                "wled": {
+                    "name": "wled",
+                    "entities": [
+                        wled_segment_1_entity.entity_id,
+                        wled_segment_2_entity.entity_id,
+                    ],
+                },
+                "hue": {
+                    "name": "hue",
+                    "entities": [
+                        hue_segment_1_entity.entity_id,
+                        hue_segment_2_entity.entity_id,
+                    ],
+                },
+                "wled_hue": {
+                    "name": "wled and hue",
+                    "entities": [
+                        wled_segment_1_entity.entity_id,
+                        wled_segment_2_entity.entity_id,
+                        hue_segment_1_entity.entity_id,
+                        hue_segment_2_entity.entity_id,
+                    ],
+                },
+            }
+        },
+    )
+
+    await async_setup_component(
+        hass,
+        "scene",
+        {
+            "scene": [
+                {
+                    "name": "scene_wled_seg_1",
+                    "entities": {wled_segment_1_entity.entity_id: "on"},
+                },
+                {
+                    "name": "scene_hue_seg_1",
+                    "entities": {hue_segment_1_entity.entity_id: "on"},
+                },
+                {
+                    "name": "scene_wled_hue",
+                    "entities": {
+                        wled_segment_1_entity.entity_id: "on",
+                        wled_segment_2_entity.entity_id: "on",
+                        hue_segment_1_entity.entity_id: "on",
+                        hue_segment_2_entity.entity_id: "on",
+                    },
+                },
+            ]
+        },
+    )
+
+    # Explore the graph from every node and make sure we find the same results
+    expected = {
+        "config_entry": {wled_config_entry.entry_id},
+        "area": {living_room_area.id},
+        "device": {wled_device.id},
+        "entity": {wled_segment_1_entity.entity_id, wled_segment_2_entity.entity_id},
+        "scene": {"scene.scene_wled_seg_1", "scene.scene_wled_hue"},
+        "group": {"group.wled", "group.wled_hue"},
+    }
+
+    for search_type, search_id in (
+        ("config_entry", wled_config_entry.entry_id),
+        ("area", living_room_area.id),
+        ("device", wled_device.id),
+        ("entity", wled_segment_1_entity.entity_id),
+        ("entity", wled_segment_2_entity.entity_id),
+        ("scene", "scene.scene_wled_seg_1"),
+        ("group", "group.wled"),
+    ):
+        searcher = search.Searcher(hass, device_reg, entity_reg)
+        results = searcher.async_search(search_type, search_id)
+        # Add the item we searched for, it's omitted from results
+        results.setdefault(search_type, set()).add(search_id)
+
+        assert (
+            results == expected
+        ), f"Results for {search_type}/{search_id} do not match up"
+
+    # For combined things, needs to return everything.
+    expected_combined = {
+        "config_entry": {wled_config_entry.entry_id, hue_config_entry.entry_id},
+        "area": {living_room_area.id, kitchen_area.id},
+        "device": {wled_device.id, hue_device.id},
+        "entity": {
+            wled_segment_1_entity.entity_id,
+            wled_segment_2_entity.entity_id,
+            hue_segment_1_entity.entity_id,
+            hue_segment_2_entity.entity_id,
+        },
+        "scene": {
+            "scene.scene_wled_seg_1",
+            "scene.scene_hue_seg_1",
+            "scene.scene_wled_hue",
+        },
+        "group": {"group.wled", "group.hue", "group.wled_hue"},
+    }
+    for search_type, search_id in (
+        ("scene", "scene.scene_wled_hue"),
+        ("group", "group.wled_hue"),
+    ):
+        searcher = search.Searcher(hass, device_reg, entity_reg)
+        results = searcher.async_search(search_type, search_id)
+        # Add the item we searched for, it's omitted from results
+        results.setdefault(search_type, set()).add(search_id)
+        assert (
+            results == expected_combined
+        ), f"Results for {search_type}/{search_id} do not match up"
+
+
+async def test_ws_api(hass, hass_ws_client):
+    """Test WS API."""
+    assert await async_setup_component(hass, "search", {})
+
+    area_reg = await hass.helpers.area_registry.async_get_registry()
+    device_reg = await hass.helpers.device_registry.async_get_registry()
+
+    kitchen_area = area_reg.async_create("Kitchen")
+
+    hue_config_entry = MockConfigEntry(domain="hue")
+    hue_config_entry.add_to_hass(hass)
+
+    hue_device = device_reg.async_get_or_create(
+        config_entry_id=hue_config_entry.entry_id,
+        name="Light Strip",
+        identifiers=({"hue", "hue-1"}),
+    )
+
+    device_reg.async_update_device(hue_device.id, area_id=kitchen_area.id)
+
+    client = await hass_ws_client(hass)
+
+    await client.send_json(
+        {
+            "id": 1,
+            "type": "search/related",
+            "item_type": "device",
+            "item_id": hue_device.id,
+        }
+    )
+    response = await client.receive_json()
+    assert response["success"]
+    assert response["result"] == {
+        "config_entry": [hue_config_entry.entry_id],
+        "area": [kitchen_area.id],
+    }