Add search integration

CODEOWNERS

@@ -279,6 +279,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/search/__init__.py

@@ -0,0 +1,132 @@
+"""The Search integration."""
+from collections import defaultdict
+
+import voluptuous as vol
+
+from homeassistant.components import websocket_api
+from homeassistant.core import HomeAssistant, 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)
+    return True
+
+
+@websocket_api.async_response
+@websocket_api.websocket_command(
+    {
+        vol.Required("type"): "search",
+        vol.Required("item_type"): vol.In(
+            ("area", "device", "entity", "script", "scene", "automation")
+        ),
+        vol.Required("item_id"): str,
+    }
+)
+async def websocket_search(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"], await searcher.search(msg["item_type"], msg["item_id"])
+    )
+
+
+class Searcher:
+    """Find related things."""
+
+    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()
+
+    async def search(self, item_type, item_id):
+        """Find results."""
+        self._add_or_resolve(item_type, item_id)
+
+        while self._to_resolve:
+            search_type, search_id = self._to_resolve.pop()
+            await getattr(self, f"_resolve_{search_type}")(search_id)
+
+        # Remove entry into graph from search results.
+        self.results[item_type].remove(item_id)
+        # Clean up entity results with types that are represented as entities
+        self.results["entity"] -= self.results["script"]
+        self.results["entity"] -= self.results["scene"]
+        self.results["entity"] -= self.results["automation"]
+
+        # Filter out empty sets.
+        return {key: val for key, val in self.results.items() if val}
+
+    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)
+        self._to_resolve.add((item_type, item_id))
+
+    async 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)
+
+    async 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)
+
+            # 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
+
+    async def _resolve_entity(self, entity_id) -> None:
+        """Resolve an entity."""
+        # Extra: Find automations, scripts, scenes that reference this 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)
+
+        domain = split_entity_id(entity_id)[0]
+
+        # We can expand these types into more types.
+        if domain in ("scene", "script", "automation"):
+            self._add_or_resolve(domain, entity_id)
+
+    async def _resolve_automation(self, automation_id) -> None:
+        """Resolve an automation."""
+        # Extra: Check with automation integration what entities/devices they reference
+
+    async def _resolve_script(self, script_id) -> None:
+        """Resolve a script."""
+        # Extra: Check with script integration what entities/devices they reference
+
+    async def _resolve_scene(self, scene_id) -> None:
+        """Resolve a scene."""
+        # Extra: Check with scene integration what entities they reference

homeassistant/components/search/manifest.json

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

tests/components/search/__init__.py

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

tests/components/search/test_init.py

@@ -0,0 +1,125 @@
+"""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_reg.async_get_or_create(
+        "light",
+        "wled",
+        "wled-1-seg-1",
+        suggested_object_id="living room light strip segment 1",
+        config_entry=wled_config_entry,
+        device_id=wled_device.id,
+    )
+    wled_segment_2 = entity_reg.async_get_or_create(
+        "light",
+        "wled",
+        "wled-1-seg-2",
+        suggested_object_id="living room light strip 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)
+
+    entity_reg.async_get_or_create(
+        "light",
+        "hue",
+        "hue-1-seg-1",
+        suggested_object_id="living room light strip segment 1",
+        config_entry=hue_config_entry,
+        device_id=hue_device.id,
+    )
+    entity_reg.async_get_or_create(
+        "light",
+        "hue",
+        "hue-1-seg-2",
+        suggested_object_id="living room light strip segment 2",
+        config_entry=hue_config_entry,
+        device_id=hue_device.id,
+    )
+
+    expected = {
+        "area": {living_room_area.id},
+        "device": {wled_device.id},
+        "entity": {wled_segment_1.entity_id, wled_segment_2.entity_id},
+    }
+
+    # Explore the graph from every node and make sure we find the same results
+    for search_type, search_id in (
+        ("area", living_room_area.id),
+        ("device", wled_device.id),
+        ("entity", wled_segment_1.entity_id),
+        ("entity", wled_segment_2.entity_id),
+    ):
+        searcher = search.Searcher(hass, device_reg, entity_reg)
+        results = await searcher.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"
+
+
+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", "item_type": "device", "item_id": hue_device.id}
+    )
+    response = await client.receive_json()
+    assert response["success"]
+    assert response["result"] == {"area": [kitchen_area.id]}