From dd2d79b77eb2ced506c22a6b3bcd4fe819a079c5 Mon Sep 17 00:00:00 2001 From: Franck Nijhof Date: Wed, 27 Mar 2024 09:39:05 +0100 Subject: [PATCH] Refactor/fix search component, including labels & floors support (#114206) Co-authored-by: TheJulianJES Co-authored-by: Robert Resch --- homeassistant/components/search/__init__.py | 649 +++++++---- tests/components/search/test_init.py | 1166 ++++++++++++------- 2 files changed, 1226 insertions(+), 589 deletions(-) diff --git a/homeassistant/components/search/__init__.py b/homeassistant/components/search/__init__.py index 1eafc137580..71b51210a25 100644 --- a/homeassistant/components/search/__init__.py +++ b/homeassistant/components/search/__init__.py @@ -2,7 +2,9 @@ from __future__ import annotations -from collections import defaultdict, deque +from collections import defaultdict +from collections.abc import Iterable +from enum import StrEnum import logging from typing import Any @@ -12,6 +14,7 @@ from homeassistant.components import automation, group, person, script, websocke from homeassistant.components.homeassistant import scene from homeassistant.core import HomeAssistant, callback, split_entity_id from homeassistant.helpers import ( + area_registry as ar, config_validation as cv, device_registry as dr, entity_registry as er, @@ -28,6 +31,25 @@ _LOGGER = logging.getLogger(__name__) CONFIG_SCHEMA = cv.empty_config_schema(DOMAIN) +# enum of item types +class ItemType(StrEnum): + """Item types.""" + + AREA = "area" + AUTOMATION = "automation" + AUTOMATION_BLUEPRINT = "automation_blueprint" + CONFIG_ENTRY = "config_entry" + DEVICE = "device" + ENTITY = "entity" + FLOOR = "floor" + GROUP = "group" + LABEL = "label" + PERSON = "person" + SCENE = "scene" + SCRIPT = "script" + SCRIPT_BLUEPRINT = "script_blueprint" + + async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: """Set up the Search component.""" websocket_api.async_register_command(hass, websocket_search_related) @@ -37,21 +59,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: @websocket_api.websocket_command( { vol.Required("type"): "search/related", - vol.Required("item_type"): vol.In( - ( - "area", - "automation", - "automation_blueprint", - "config_entry", - "device", - "entity", - "group", - "person", - "scene", - "script", - "script_blueprint", - ) - ), + vol.Required("item_type"): vol.Coerce(ItemType), vol.Required("item_id"): str, } ) @@ -62,271 +70,520 @@ def websocket_search_related( msg: dict[str, Any], ) -> None: """Handle search.""" - searcher = Searcher( - hass, - dr.async_get(hass), - er.async_get(hass), - get_entity_sources(hass), - ) + searcher = Searcher(hass, get_entity_sources(hass)) connection.send_result( msg["id"], searcher.async_search(msg["item_type"], msg["item_id"]) ) class Searcher: - """Find related things. + """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 = { - "area", - "automation", - "automation_blueprint", - "config_entry", - "group", - "scene", - "script", - "script_blueprint", - } - # These types exist as an entity and so need cleanup in results EXIST_AS_ENTITY = {"automation", "group", "person", "scene", "script"} def __init__( self, hass: HomeAssistant, - device_reg: dr.DeviceRegistry, - entity_reg: er.EntityRegistry, entity_sources: dict[str, EntityInfo], ) -> None: """Search results.""" self.hass = hass - self._device_reg = device_reg - self._entity_reg = entity_reg - self._sources = entity_sources - self.results: defaultdict[str, set[str]] = defaultdict(set) - self._to_resolve: deque[tuple[str, str]] = deque() + self._area_registry = ar.async_get(hass) + self._device_registry = dr.async_get(hass) + self._entity_registry = er.async_get(hass) + self._entity_sources = entity_sources + self.results: defaultdict[ItemType, set[str]] = defaultdict(set) @callback - def async_search(self, item_type: str, item_id: str) -> dict[str, set[str]]: + def async_search(self, item_type: ItemType, item_id: str) -> dict[str, set[str]]: """Find results.""" _LOGGER.debug("Searching for %s/%s", item_type, item_id) - self.results[item_type].add(item_id) - self._to_resolve.append((item_type, item_id)) + getattr(self, f"_async_search_{item_type}")(item_id) - while self._to_resolve: - search_type, search_id = self._to_resolve.popleft() - 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. - for result_type in self.EXIST_AS_ENTITY: - self.results["entity"] -= self.results[result_type] - - # Remove entry into graph from search results. - to_remove_item_type = item_type - if item_type == "entity": - domain = split_entity_id(item_id)[0] - - if domain in self.EXIST_AS_ENTITY: - to_remove_item_type = domain - - self.results[to_remove_item_type].remove(item_id) + # Remove the original requested item from the results (if present) + if item_type in self.results and item_id in self.results[item_type]: + 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: str, item_id: str) -> None: - """Add an item to explore.""" - if item_id in self.results[item_type]: + def _add(self, item_type: ItemType, item_id: str | Iterable[str] | None) -> None: + """Add an item (or items) to the results.""" + if item_id is None: return - self.results[item_type].add(item_id) - - if item_type not in self.DONT_RESOLVE: - self._to_resolve.append((item_type, item_id)) + if isinstance(item_id, str): + self.results[item_type].add(item_id) + else: + self.results[item_type].update(item_id) @callback - def _resolve_area(self, area_id: str) -> None: - """Resolve an area.""" - for device in dr.async_entries_for_area(self._device_reg, area_id): - self._add_or_resolve("device", device.id) + def _async_search_area(self, area_id: str, *, entry_point: bool = True) -> None: + """Find results for an area.""" + if not (area_entry := self._async_resolve_up_area(area_id)): + return - for entity_entry in er.async_entries_for_area(self._entity_reg, area_id): - self._add_or_resolve("entity", entity_entry.entity_id) + if entry_point: + # Add labels of this area + self._add(ItemType.LABEL, area_entry.labels) - for entity_id in script.scripts_with_area(self.hass, area_id): - self._add_or_resolve("entity", entity_id) + # Automations referencing this area + self._add( + ItemType.AUTOMATION, automation.automations_with_area(self.hass, area_id) + ) - for entity_id in automation.automations_with_area(self.hass, area_id): - self._add_or_resolve("entity", entity_id) + # Scripts referencing this area + self._add(ItemType.SCRIPT, script.scripts_with_area(self.hass, 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) + + # Config entries for devices in this area + if device_entry := self._device_registry.async_get(device.id): + self._add(ItemType.CONFIG_ENTRY, device_entry.config_entries) + + # Automations referencing this device + self._add( + ItemType.AUTOMATION, + automation.automations_with_device(self.hass, device.id), + ) + + # Scripts referencing this device + self._add(ItemType.SCRIPT, script.scripts_with_device(self.hass, device.id)) + + # Entities of this device + for entity_entry in er.async_entries_for_device( + self._entity_registry, device.id + ): + # 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) + + # Entities in this area + for entity_entry in er.async_entries_for_area(self._entity_registry, area_id): + self._add(ItemType.ENTITY, entity_entry.entity_id) + + # If this entity also exists as a resource, we add it. + if entity_entry.domain in self.EXIST_AS_ENTITY: + self._add(ItemType(entity_entry.domain), entity_entry.entity_id) + + # Automations referencing this entity + self._add( + ItemType.AUTOMATION, + automation.automations_with_entity(self.hass, entity_entry.entity_id), + ) + + # Scripts referencing this entity + self._add( + ItemType.SCRIPT, + script.scripts_with_entity(self.hass, entity_entry.entity_id), + ) + + # Groups that have this entity as a member + self._add( + ItemType.GROUP, + group.groups_with_entity(self.hass, entity_entry.entity_id), + ) + + # Persons that use this entity + self._add( + ItemType.PERSON, + person.persons_with_entity(self.hass, entity_entry.entity_id), + ) + + # Scenes that reference this entity + self._add( + ItemType.SCENE, + scene.scenes_with_entity(self.hass, entity_entry.entity_id), + ) + + # Config entries for entities in this area + self._add(ItemType.CONFIG_ENTRY, entity_entry.config_entry_id) @callback - def _resolve_automation(self, automation_entity_id: str) -> None: - """Resolve an automation. + def _async_search_automation(self, automation_entity_id: str) -> None: + """Find results for an automation.""" + # Up resolve the automation entity itself + if entity_entry := self._async_resolve_up_entity(automation_entity_id): + # Add labels of this automation entity + self._add(ItemType.LABEL, entity_entry.labels) - Will only be called if automation is an entry point. - """ - for entity in automation.entities_in_automation( - self.hass, automation_entity_id - ): - self._add_or_resolve("entity", entity) + # Find the blueprint used in this automation + self._add( + ItemType.AUTOMATION_BLUEPRINT, + automation.blueprint_in_automation(self.hass, automation_entity_id), + ) - for device in automation.devices_in_automation(self.hass, automation_entity_id): - self._add_or_resolve("device", device) + # Floors referenced in this automation + self._add( + ItemType.FLOOR, + automation.floors_in_automation(self.hass, automation_entity_id), + ) + # Areas referenced in this automation for area in automation.areas_in_automation(self.hass, automation_entity_id): - self._add_or_resolve("area", area) + self._add(ItemType.AREA, area) + self._async_resolve_up_area(area) - if blueprint := automation.blueprint_in_automation( + # Devices referenced in this automation + for device in automation.devices_in_automation(self.hass, automation_entity_id): + self._add(ItemType.DEVICE, device) + self._async_resolve_up_device(device) + + # Entities referenced in this automation + for entity_id in automation.entities_in_automation( self.hass, automation_entity_id ): - self._add_or_resolve("automation_blueprint", blueprint) + self._add(ItemType.ENTITY, entity_id) + self._async_resolve_up_entity(entity_id) + + # If this entity also exists as a resource, we add it. + domain = split_entity_id(entity_id)[0] + if domain in self.EXIST_AS_ENTITY: + self._add(ItemType(domain), entity_id) + + # For an automation, we want to unwrap the groups, to ensure we + # relate this automation to all those members as well. + if domain == "group": + for group_entity_id in group.get_entity_ids(self.hass, entity_id): + self._add(ItemType.ENTITY, group_entity_id) + self._async_resolve_up_entity(group_entity_id) + + # For an automation, we want to unwrap the scenes, to ensure we + # relate this automation to all referenced entities as well. + if domain == "scene": + for scene_entity_id in scene.entities_in_scene(self.hass, entity_id): + self._add(ItemType.ENTITY, scene_entity_id) + self._async_resolve_up_entity(scene_entity_id) + + # Fully search the script if it is part of an automation. + # This makes the automation return all results of the embedded script. + if domain == "script": + self._async_search_script(entity_id, entry_point=False) @callback - def _resolve_automation_blueprint(self, blueprint_path: str) -> None: - """Resolve an automation blueprint. - - Will only be called if blueprint is an entry point. - """ - for entity_id in automation.automations_with_blueprint( - self.hass, blueprint_path - ): - self._add_or_resolve("automation", entity_id) + def _async_search_automation_blueprint(self, blueprint_path: str) -> None: + """Find results for an automation blueprint.""" + self._add( + ItemType.AUTOMATION, + automation.automations_with_blueprint(self.hass, blueprint_path), + ) @callback - def _resolve_config_entry(self, config_entry_id: str) -> None: - """Resolve a config entry. - - Will only be called if config entry is an entry point. - """ + def _async_search_config_entry(self, config_entry_id: str) -> None: + """Find results for a config entry.""" for device_entry in dr.async_entries_for_config_entry( - self._device_reg, config_entry_id + self._device_registry, config_entry_id ): - self._add_or_resolve("device", device_entry.id) + self._add(ItemType.DEVICE, device_entry.id) + self._async_search_device(device_entry.id, entry_point=False) for entity_entry in er.async_entries_for_config_entry( - self._entity_reg, config_entry_id + self._entity_registry, config_entry_id ): - self._add_or_resolve("entity", entity_entry.entity_id) + self._add(ItemType.ENTITY, entity_entry.entity_id) + self._async_search_entity(entity_entry.entity_id, entry_point=False) @callback - def _resolve_device(self, device_id: str) -> 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) + def _async_search_device(self, device_id: str, *, entry_point: bool = True) -> None: + """Find results for a device.""" + if not (device_entry := self._async_resolve_up_device(device_id)): + return - for config_entry_id in device_entry.config_entries: - self._add_or_resolve("config_entry", config_entry_id) + if entry_point: + # Add labels of this device + self._add(ItemType.LABEL, device_entry.labels) - # We do not resolve device_entry.via_device_id because that - # device is not related data-wise inside HA. + # Automations referencing this device + self._add( + ItemType.AUTOMATION, + automation.automations_with_device(self.hass, device_id), + ) - for entity_entry in er.async_entries_for_device(self._entity_reg, device_id): - self._add_or_resolve("entity", entity_entry.entity_id) + # Scripts referencing this device + self._add(ItemType.SCRIPT, script.scripts_with_device(self.hass, device_id)) - for entity_id in script.scripts_with_device(self.hass, device_id): - self._add_or_resolve("entity", entity_id) - - for entity_id in automation.automations_with_device(self.hass, device_id): - self._add_or_resolve("entity", entity_id) + # Entities of this device + for entity_entry in er.async_entries_for_device( + self._entity_registry, device_id + ): + self._add(ItemType.ENTITY, entity_entry.entity_id) + # Add all entity information as well + self._async_search_entity(entity_entry.entity_id, entry_point=False) @callback - def _resolve_entity(self, entity_id: str) -> None: - """Resolve an entity.""" - # Extra: Find automations and scripts that reference this entity. + def _async_search_entity(self, entity_id: str, *, entry_point: bool = True) -> None: + """Find results for an entity.""" + # Resolve up the entity itself + entity_entry = self._async_resolve_up_entity(entity_id) - for entity in scene.scenes_with_entity(self.hass, entity_id): - self._add_or_resolve("entity", entity) + if entity_entry and entry_point: + # Add labels of this entity + self._add(ItemType.LABEL, entity_entry.labels) - for entity in group.groups_with_entity(self.hass, entity_id): - self._add_or_resolve("entity", entity) + # Automations referencing this entity + self._add( + ItemType.AUTOMATION, + automation.automations_with_entity(self.hass, entity_id), + ) - for entity in automation.automations_with_entity(self.hass, entity_id): - self._add_or_resolve("entity", entity) + # Scripts referencing this entity + self._add(ItemType.SCRIPT, script.scripts_with_entity(self.hass, entity_id)) - for entity in script.scripts_with_entity(self.hass, entity_id): - self._add_or_resolve("entity", entity) + # Groups that have this entity as a member + self._add(ItemType.GROUP, group.groups_with_entity(self.hass, entity_id)) - for entity in person.persons_with_entity(self.hass, entity_id): - self._add_or_resolve("entity", entity) + # Persons referencing this entity + self._add(ItemType.PERSON, person.persons_with_entity(self.hass, entity_id)) - # 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) - else: - source = self._sources.get(entity_id) - if source is not None and "config_entry" in source: - self._add_or_resolve("config_entry", source["config_entry"]) - - domain = split_entity_id(entity_id)[0] - - if domain in self.EXIST_AS_ENTITY: - self._add_or_resolve(domain, entity_id) + # Scenes referencing this entity + self._add(ItemType.SCENE, scene.scenes_with_entity(self.hass, entity_id)) @callback - def _resolve_group(self, group_entity_id: str) -> None: - """Resolve a group. + def _async_search_floor(self, floor_id: str) -> None: + """Find results for a floor.""" + # Automations referencing this floor + self._add( + ItemType.AUTOMATION, + automation.automations_with_floor(self.hass, floor_id), + ) - Will only be called if group is an entry point. + # Scripts referencing this floor + self._add(ItemType.SCRIPT, script.scripts_with_floor(self.hass, floor_id)) + + for area_entry in ar.async_entries_for_floor(self._area_registry, floor_id): + self._add(ItemType.AREA, area_entry.id) + self._async_search_area(area_entry.id, entry_point=False) + + @callback + def _async_search_group(self, group_entity_id: str) -> None: + """Find results for a group. + + Note: We currently only support the classic groups, thus + we don't look up the area/floor for a group entity. """ + # Automations referencing this group + self._add( + ItemType.AUTOMATION, + automation.automations_with_entity(self.hass, group_entity_id), + ) + + # Scripts referencing this group + self._add( + ItemType.SCRIPT, script.scripts_with_entity(self.hass, group_entity_id) + ) + + # Scenes that reference this group + self._add(ItemType.SCENE, scene.scenes_with_entity(self.hass, group_entity_id)) + + # Entities in this group for entity_id in group.get_entity_ids(self.hass, group_entity_id): - self._add_or_resolve("entity", entity_id) + self._add(ItemType.ENTITY, entity_id) + self._async_resolve_up_entity(entity_id) @callback - def _resolve_person(self, person_entity_id: str) -> None: - """Resolve a person. + def _async_search_label(self, label_id: str) -> None: + """Find results for a label.""" - Will only be called if person is an entry point. - """ - for entity in person.entities_in_person(self.hass, person_entity_id): - self._add_or_resolve("entity", entity) + # Areas with this label + for area_entry in ar.async_entries_for_label(self._area_registry, label_id): + self._add(ItemType.AREA, area_entry.id) + + # Devices with this label + for device in dr.async_entries_for_label(self._device_registry, label_id): + self._add(ItemType.DEVICE, device.id) + + # Entities with this label + for entity_entry in er.async_entries_for_label(self._entity_registry, label_id): + self._add(ItemType.ENTITY, entity_entry.entity_id) + + # If this entity also exists as a resource, we add it. + domain = split_entity_id(entity_entry.entity_id)[0] + if domain in self.EXIST_AS_ENTITY: + self._add(ItemType(domain), entity_entry.entity_id) + + # Automations referencing this label + self._add( + ItemType.AUTOMATION, + automation.automations_with_label(self.hass, label_id), + ) + + # Scripts referencing this label + self._add(ItemType.SCRIPT, script.scripts_with_label(self.hass, label_id)) @callback - def _resolve_scene(self, scene_entity_id: str) -> None: - """Resolve a scene. + def _async_search_person(self, person_entity_id: str) -> None: + """Find results for a person.""" + # Up resolve the scene entity itself + if entity_entry := self._async_resolve_up_entity(person_entity_id): + # Add labels of this person entity + self._add(ItemType.LABEL, entity_entry.labels) - Will only be called if scene is an entry point. - """ + # Automations referencing this person + self._add( + ItemType.AUTOMATION, + automation.automations_with_entity(self.hass, person_entity_id), + ) + + # Scripts referencing this person + self._add( + ItemType.SCRIPT, script.scripts_with_entity(self.hass, person_entity_id) + ) + + # Add all member entities of this person + self._add( + ItemType.ENTITY, person.entities_in_person(self.hass, person_entity_id) + ) + + @callback + def _async_search_scene(self, scene_entity_id: str) -> None: + """Find results for a scene.""" + # Up resolve the scene entity itself + if entity_entry := self._async_resolve_up_entity(scene_entity_id): + # Add labels of this scene entity + self._add(ItemType.LABEL, entity_entry.labels) + + # Automations referencing this scene + self._add( + ItemType.AUTOMATION, + automation.automations_with_entity(self.hass, scene_entity_id), + ) + + # Scripts referencing this scene + self._add( + ItemType.SCRIPT, script.scripts_with_entity(self.hass, scene_entity_id) + ) + + # Add all entities in this scene for entity in scene.entities_in_scene(self.hass, scene_entity_id): - self._add_or_resolve("entity", entity) + self._add(ItemType.ENTITY, entity) + self._async_resolve_up_entity(entity) @callback - def _resolve_script(self, script_entity_id: str) -> None: - """Resolve a script. + def _async_search_script( + self, script_entity_id: str, *, entry_point: bool = True + ) -> None: + """Find results for a script.""" + # Up resolve the script entity itself + entity_entry = self._async_resolve_up_entity(script_entity_id) - Will only be called if script is an entry point. - """ - for entity in script.entities_in_script(self.hass, script_entity_id): - self._add_or_resolve("entity", entity) + if entity_entry and entry_point: + # Add labels of this script entity + self._add(ItemType.LABEL, entity_entry.labels) - for device in script.devices_in_script(self.hass, script_entity_id): - self._add_or_resolve("device", device) + # Find the blueprint used in this script + self._add( + ItemType.SCRIPT_BLUEPRINT, + script.blueprint_in_script(self.hass, script_entity_id), + ) + # Floors referenced in this script + self._add(ItemType.FLOOR, script.floors_in_script(self.hass, script_entity_id)) + + # Areas referenced in this script for area in script.areas_in_script(self.hass, script_entity_id): - self._add_or_resolve("area", area) + self._add(ItemType.AREA, area) + self._async_resolve_up_area(area) - if blueprint := script.blueprint_in_script(self.hass, script_entity_id): - self._add_or_resolve("script_blueprint", blueprint) + # Devices referenced in this script + for device in script.devices_in_script(self.hass, script_entity_id): + self._add(ItemType.DEVICE, device) + self._async_resolve_up_device(device) + + # Entities referenced in this script + for entity_id in script.entities_in_script(self.hass, script_entity_id): + self._add(ItemType.ENTITY, entity_id) + self._async_resolve_up_entity(entity_id) + + # If this entity also exists as a resource, we add it. + domain = split_entity_id(entity_id)[0] + if domain in self.EXIST_AS_ENTITY: + self._add(ItemType(domain), entity_id) + + # For an script, we want to unwrap the groups, to ensure we + # relate this script to all those members as well. + if domain == "group": + for group_entity_id in group.get_entity_ids(self.hass, entity_id): + self._add(ItemType.ENTITY, group_entity_id) + self._async_resolve_up_entity(group_entity_id) + + # For an script, we want to unwrap the scenes, to ensure we + # relate this script to all referenced entities as well. + if domain == "scene": + for scene_entity_id in scene.entities_in_scene(self.hass, entity_id): + self._add(ItemType.ENTITY, scene_entity_id) + self._async_resolve_up_entity(scene_entity_id) + + # Fully search the script if it is nested. + # This makes the script return all results of the embedded script. + if domain == "script": + self._async_search_script(entity_id, entry_point=False) @callback - def _resolve_script_blueprint(self, blueprint_path: str) -> None: - """Resolve a script blueprint. + def _async_search_script_blueprint(self, blueprint_path: str) -> None: + """Find results for a script blueprint.""" + self._add( + ItemType.SCRIPT, script.scripts_with_blueprint(self.hass, blueprint_path) + ) - Will only be called if blueprint is an entry point. + @callback + def _async_resolve_up_device(self, device_id: str) -> dr.DeviceEntry | None: + """Resolve up from a device. + + Above a device is an area or floor. + Above a device is also the config entry. """ - for entity_id in script.scripts_with_blueprint(self.hass, blueprint_path): - self._add_or_resolve("script", entity_id) + if device_entry := self._device_registry.async_get(device_id): + if device_entry.area_id: + self._add(ItemType.AREA, device_entry.area_id) + self._async_resolve_up_area(device_entry.area_id) + + self._add(ItemType.CONFIG_ENTRY, device_entry.config_entries) + + return device_entry + + @callback + def _async_resolve_up_entity(self, entity_id: str) -> er.RegistryEntry | None: + """Resolve up from an entity. + + Above an entity is a device, area or floor. + Above an entity is also the config entry. + """ + if entity_entry := self._entity_registry.async_get(entity_id): + # Entity has an overridden area + if entity_entry.area_id: + self._add(ItemType.AREA, entity_entry.area_id) + self._async_resolve_up_area(entity_entry.area_id) + + # Inherit area from device + elif entity_entry.device_id and ( + device_entry := self._device_registry.async_get(entity_entry.device_id) + ): + if device_entry.area_id: + self._add(ItemType.AREA, device_entry.area_id) + self._async_resolve_up_area(device_entry.area_id) + + # Add device that provided this entity + self._add(ItemType.DEVICE, entity_entry.device_id) + + # Add config entry that provided this entity + self._add(ItemType.CONFIG_ENTRY, entity_entry.config_entry_id) + elif source := self._entity_sources.get(entity_id): + # Add config entry that provided this entity + self._add(ItemType.CONFIG_ENTRY, source.get("config_entry")) + + return entity_entry + + @callback + def _async_resolve_up_area(self, area_id: str) -> ar.AreaEntry | None: + """Resolve up from an area. + + Above an area can be a floor. + """ + if area_entry := self._area_registry.async_get_area(area_id): + self._add(ItemType.FLOOR, area_entry.floor_id) + + return area_entry diff --git a/tests/components/search/test_init.py b/tests/components/search/test_init.py index 57b72de2bcd..ee7b60dc9ac 100644 --- a/tests/components/search/test_init.py +++ b/tests/components/search/test_init.py @@ -1,14 +1,18 @@ """Tests for Search integration.""" import pytest +from pytest_unordered import unordered -from homeassistant.components import search +from homeassistant.components.search import ItemType, Searcher from homeassistant.core import HomeAssistant from homeassistant.helpers import ( area_registry as ar, device_registry as dr, entity_registry as er, + floor_registry as fr, + label_registry as lr, ) +from homeassistant.helpers.entity import EntityInfo from homeassistant.setup import async_setup_component from tests.common import MockConfigEntry @@ -20,38 +24,79 @@ def stub_blueprint_populate_autouse(stub_blueprint_populate: None) -> None: """Stub copying the blueprints to the config folder.""" -MOCK_ENTITY_SOURCES = { - "light.platform_config_source": { - "domain": "wled", - }, - "light.config_entry_source": { - "config_entry": "config_entry_id", - "domain": "wled", - }, -} - - async def test_search( hass: HomeAssistant, area_registry: ar.AreaRegistry, device_registry: dr.DeviceRegistry, entity_registry: er.EntityRegistry, + floor_registry: fr.FloorRegistry, + label_registry: lr.LabelRegistry, + hass_ws_client: WebSocketGenerator, ) -> None: - """Test that search works.""" - living_room_area = area_registry.async_create("Living Room") + """Test search.""" + assert await async_setup_component(hass, "search", {}) - # Light strip with 2 lights. + # Labels + label_energy = label_registry.async_create("Energy") + label_christmas = label_registry.async_create("Christmas") + label_other = label_registry.async_create("Other") + + # Floors + first_floor = floor_registry.async_create("First Floor") + second_floor = floor_registry.async_create("Second Floor") + + # Areas + bedroom_area = area_registry.async_create( + "Bedroom", floor_id=second_floor.floor_id, labels={label_other.label_id} + ) + kitchen_area = area_registry.async_create("Kitchen", floor_id=first_floor.floor_id) + living_room_area = area_registry.async_create( + "Living Room", floor_id=first_floor.floor_id + ) + + # Config entries + hue_config_entry = MockConfigEntry(domain="hue") + hue_config_entry.add_to_hass(hass) wled_config_entry = MockConfigEntry(domain="wled") wled_config_entry.add_to_hass(hass) + # Devices + hue_device = device_registry.async_get_or_create( + config_entry_id=hue_config_entry.entry_id, + name="Light Strip", + identifiers={("hue", "hue-1")}, + ) + device_registry.async_update_device(hue_device.id, area_id=kitchen_area.id) + wled_device = device_registry.async_get_or_create( config_entry_id=wled_config_entry.entry_id, name="Light Strip", identifiers=({"wled", "wled-1"}), ) + device_registry.async_update_device( + wled_device.id, area_id=living_room_area.id, labels={label_christmas.label_id} + ) - device_registry.async_update_device(wled_device.id, area_id=living_room_area.id) - + # Entities + hue_segment_1_entity = entity_registry.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, + ) + entity_registry.async_update_entity( + hue_segment_1_entity.entity_id, labels={label_energy.label_id} + ) + hue_segment_2_entity = entity_registry.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, + ) wled_segment_1_entity = entity_registry.async_get_or_create( "light", "wled", @@ -68,48 +113,59 @@ async def test_search( config_entry=wled_config_entry, device_id=wled_device.id, ) + entity_registry.async_update_entity( + wled_segment_2_entity.entity_id, area_id=bedroom_area.id + ) + scene_wled_hue_entity = entity_registry.async_get_or_create( + "scene", + "homeassistant", + "wled_hue", + suggested_object_id="scene_wled_hue", + ) + entity_registry.async_update_entity( + scene_wled_hue_entity.entity_id, + area_id=bedroom_area.id, + labels={label_other.label_id}, + ) + + # Persons can technically be assigned to areas + person_paulus_entity = entity_registry.async_get_or_create( + "person", + "person", + "abcd", + suggested_object_id="paulus", + ) + entity_registry.async_update_entity( + person_paulus_entity.entity_id, + area_id=bedroom_area.id, + labels={label_other.label_id}, + ) + + script_scene_entity = entity_registry.async_get_or_create( + "script", + "script", + "scene", + suggested_object_id="scene", + ) + entity_registry.async_update_entity( + script_scene_entity.entity_id, + area_id=bedroom_area.id, + labels={label_other.label_id}, + ) + + # Entity sources entity_sources = { - "light.wled_platform_config_source": { - "domain": "wled", - }, - "light.wled_config_entry_source": { - "config_entry": wled_config_entry.entry_id, - "domain": "wled", - }, + "light.wled_platform_config_source": EntityInfo( + domain="wled", + ), + "light.wled_config_entry_source": EntityInfo( + config_entry=wled_config_entry.entry_id, + domain="wled", + ), } - # Non related info. - kitchen_area = area_registry.async_create("Kitchen") - - hue_config_entry = MockConfigEntry(domain="hue") - hue_config_entry.add_to_hass(hass) - - hue_device = device_registry.async_get_or_create( - config_entry_id=hue_config_entry.entry_id, - name="Light Strip", - identifiers=({"hue", "hue-1"}), - ) - - device_registry.async_update_device(hue_device.id, area_id=kitchen_area.id) - - hue_segment_1_entity = entity_registry.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_registry.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, - ) - + # Groups await async_setup_component( hass, "group", @@ -142,6 +198,22 @@ async def test_search( }, ) + # Persons + assert await async_setup_component( + hass, + "person", + { + "person": [ + { + "id": "abcd", + "name": "Paulus", + "device_trackers": ["device_tracker.paulus_iphone"], + } + ] + }, + ) + + # Scenes await async_setup_component( hass, "scene", @@ -156,6 +228,7 @@ async def test_search( "entities": {hue_segment_1_entity.entity_id: "on"}, }, { + "id": "wled_hue", "name": "scene_wled_hue", "entities": { wled_segment_1_entity.entity_id: "on", @@ -168,11 +241,144 @@ async def test_search( }, ) - await async_setup_component( + # Automations + assert await async_setup_component( + hass, + "automation", + { + "automation": [ + { + "id": "unique_id", + "alias": "blueprint_automation_1", + "trigger": {"platform": "template", "value_template": "true"}, + "use_blueprint": { + "path": "test_event_service.yaml", + "input": { + "trigger_event": "blueprint_event_1", + "service_to_call": "test.automation_1", + "a_number": 5, + }, + }, + }, + { + "alias": "blueprint_automation_2", + "trigger": {"platform": "template", "value_template": "true"}, + "use_blueprint": { + "path": "test_event_service.yaml", + "input": { + "trigger_event": "blueprint_event_2", + "service_to_call": "test.automation_2", + "a_number": 5, + }, + }, + }, + { + "alias": "wled_entity", + "trigger": {"platform": "template", "value_template": "true"}, + "action": [ + { + "service": "test.script", + "data": {"entity_id": wled_segment_1_entity.entity_id}, + }, + ], + }, + { + "alias": "wled_device", + "trigger": {"platform": "template", "value_template": "true"}, + "action": [ + { + "domain": "light", + "device_id": wled_device.id, + "entity_id": wled_segment_1_entity.entity_id, + "type": "turn_on", + }, + ], + }, + { + "alias": "floor", + "trigger": {"platform": "template", "value_template": "true"}, + "action": [ + { + "service": "test.script", + "target": {"floor_id": first_floor.floor_id}, + }, + ], + }, + { + "alias": "area", + "trigger": {"platform": "template", "value_template": "true"}, + "action": [ + { + "service": "test.script", + "target": {"area_id": kitchen_area.id}, + }, + ], + }, + { + "alias": "group", + "trigger": {"platform": "template", "value_template": "true"}, + "action": [ + { + "service": "homeassistant.turn_on", + "target": {"entity_id": "group.wled_hue"}, + }, + ], + }, + { + "alias": "scene", + "trigger": {"platform": "template", "value_template": "true"}, + "action": [ + { + "scene": scene_wled_hue_entity.entity_id, + }, + ], + }, + { + "alias": "script", + "trigger": {"platform": "template", "value_template": "true"}, + "action": [ + { + "service": "script.turn_on", + "data": {"entity_id": script_scene_entity.entity_id}, + }, + ], + }, + { + "alias": "label", + "trigger": {"platform": "template", "value_template": "true"}, + "action": [ + { + "service": "script.turn_on", + "target": {"label_id": label_christmas.label_id}, + }, + ], + }, + ] + }, + ) + + # Scripts + assert await async_setup_component( hass, "script", { "script": { + "blueprint_script_1": { + "use_blueprint": { + "path": "test_service.yaml", + "input": { + "service_to_call": "test.automation", + }, + } + }, + "blueprint_script_2": { + "use_blueprint": { + "path": "test_service.yaml", + "input": { + "service_to_call": "test.automation", + }, + } + }, "wled": { "sequence": [ { @@ -205,375 +411,532 @@ async def test_search( }, ] }, + "device": { + "sequence": [ + { + "service": "test.script", + "target": {"device_id": hue_device.id}, + }, + ], + }, + "floor": { + "sequence": [ + { + "service": "test.script", + "target": {"floor_id": first_floor.floor_id}, + }, + ], + }, + "area": { + "sequence": [ + { + "service": "test.script", + "target": {"area_id": kitchen_area.id}, + }, + ], + }, + "group": { + "sequence": [ + { + "service": "test.script", + "target": {"entity_id": "group.wled_hue"}, + }, + ], + }, + "scene": { + "sequence": [ + { + "scene": scene_wled_hue_entity.entity_id, + }, + ], + }, + "label": { + "sequence": [ + { + "service": "test.script", + "target": {"label_id": label_other.label_id}, + }, + ], + }, + "nested": { + "sequence": [ + { + "service": "script.turn_on", + "data": {"entity_id": script_scene_entity.entity_id}, + }, + ], + }, } }, ) - assert await async_setup_component( - hass, - "automation", - { - "automation": [ - { - "alias": "wled_entity", - "trigger": {"platform": "template", "value_template": "true"}, - "action": [ - { - "service": "test.script", - "data": {"entity_id": wled_segment_1_entity.entity_id}, - }, - ], - }, - { - "alias": "wled_device", - "trigger": {"platform": "template", "value_template": "true"}, - "action": [ - { - "domain": "light", - "device_id": wled_device.id, - "entity_id": wled_segment_1_entity.entity_id, - "type": "turn_on", - }, - ], - }, - ] + def search(item_type: ItemType, item_id: str) -> dict[str, set[str]]: + """Search.""" + searcher = Searcher(hass, entity_sources) + return searcher.async_search(item_type, item_id) + + # + # Tests + # + assert not search(ItemType.AREA, "unknown") + assert search(ItemType.AREA, bedroom_area.id) == { + ItemType.AUTOMATION: {"automation.scene", "automation.script"}, + ItemType.CONFIG_ENTRY: {wled_config_entry.entry_id}, + ItemType.ENTITY: { + wled_segment_2_entity.entity_id, + scene_wled_hue_entity.entity_id, + script_scene_entity.entity_id, + person_paulus_entity.entity_id, }, - ) - - # Ensure automations set up correctly. - assert hass.states.get("automation.wled_entity") is not None - assert hass.states.get("automation.wled_device") is not None - - # 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"}, - "script": {"script.wled"}, - "automation": {"automation.wled_entity", "automation.wled_device"}, + ItemType.FLOOR: {second_floor.floor_id}, + ItemType.GROUP: {"group.wled", "group.wled_hue"}, + ItemType.LABEL: {label_other.label_id}, + ItemType.PERSON: {person_paulus_entity.entity_id}, + ItemType.SCENE: {scene_wled_hue_entity.entity_id}, + ItemType.SCRIPT: {script_scene_entity.entity_id, "script.nested"}, + } + assert search(ItemType.AREA, living_room_area.id) == { + ItemType.AUTOMATION: {"automation.wled_device"}, + 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}, + } + assert search(ItemType.AREA, kitchen_area.id) == { + ItemType.AUTOMATION: {"automation.area"}, + ItemType.CONFIG_ENTRY: {hue_config_entry.entry_id}, + ItemType.DEVICE: {hue_device.id}, + ItemType.ENTITY: { + hue_segment_1_entity.entity_id, + hue_segment_2_entity.entity_id, + }, + ItemType.FLOOR: {first_floor.floor_id}, + ItemType.SCRIPT: {"script.area", "script.device"}, } - 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"), - ("script", "script.wled"), - ("automation", "automation.wled_entity"), - ("automation", "automation.wled_device"), - ): - searcher = search.Searcher( - hass, device_registry, entity_registry, entity_sources - ) - 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": { + assert not search(ItemType.AUTOMATION, "automation.unknown") + assert search(ItemType.AUTOMATION, "automation.blueprint_automation_1") == { + ItemType.AUTOMATION_BLUEPRINT: {"test_event_service.yaml"}, + ItemType.ENTITY: {"light.kitchen"}, + } + assert search(ItemType.AUTOMATION, "automation.blueprint_automation_2") == { + ItemType.AUTOMATION_BLUEPRINT: {"test_event_service.yaml"}, + ItemType.ENTITY: {"light.kitchen"}, + } + assert search(ItemType.AUTOMATION, "automation.wled_entity") == { + ItemType.AREA: {living_room_area.id}, + 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}, + } + assert search(ItemType.AUTOMATION, "automation.wled_device") == { + ItemType.AREA: {living_room_area.id}, + ItemType.CONFIG_ENTRY: {wled_config_entry.entry_id}, + ItemType.DEVICE: {wled_device.id}, + ItemType.FLOOR: {first_floor.floor_id}, + } + assert search(ItemType.AUTOMATION, "automation.floor") == { + ItemType.FLOOR: {first_floor.floor_id}, + } + assert search(ItemType.AUTOMATION, "automation.area") == { + ItemType.AREA: {kitchen_area.id}, + ItemType.FLOOR: {first_floor.floor_id}, + } + assert search(ItemType.AUTOMATION, "automation.group") == { + ItemType.AREA: {bedroom_area.id, living_room_area.id, kitchen_area.id}, + ItemType.CONFIG_ENTRY: {wled_config_entry.entry_id, hue_config_entry.entry_id}, + ItemType.DEVICE: {wled_device.id, hue_device.id}, + ItemType.ENTITY: { + "group.wled_hue", 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", + ItemType.FLOOR: {first_floor.floor_id, second_floor.floor_id}, + ItemType.GROUP: {"group.wled_hue"}, + } + assert search(ItemType.AUTOMATION, "automation.scene") == { + ItemType.AREA: {bedroom_area.id, kitchen_area.id, living_room_area.id}, + ItemType.CONFIG_ENTRY: {wled_config_entry.entry_id, hue_config_entry.entry_id}, + ItemType.DEVICE: {wled_device.id, hue_device.id}, + ItemType.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_wled_hue_entity.entity_id, }, - "group": {"group.wled", "group.hue", "group.wled_hue"}, - "script": {"script.wled", "script.hue"}, - "automation": {"automation.wled_entity", "automation.wled_device"}, + ItemType.FLOOR: {first_floor.floor_id, second_floor.floor_id}, + ItemType.SCENE: {scene_wled_hue_entity.entity_id}, } - for search_type, search_id in ( - ("scene", "scene.scene_wled_hue"), - ("group", "group.wled_hue"), - ): - searcher = search.Searcher( - hass, device_registry, entity_registry, entity_sources - ) - 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" - - for search_type, search_id in ( - ("entity", "automation.non_existing"), - ("entity", "scene.non_existing"), - ("entity", "group.non_existing"), - ("entity", "script.non_existing"), - ("entity", "light.non_existing"), - ("area", "non_existing"), - ("config_entry", "non_existing"), - ("device", "non_existing"), - ("group", "group.non_existing"), - ("scene", "scene.non_existing"), - ("script", "script.non_existing"), - ("automation", "automation.non_existing"), - ): - searcher = search.Searcher( - hass, device_registry, entity_registry, entity_sources - ) - assert searcher.async_search(search_type, search_id) == {} - - # Test search of templated script. We can't find referenced areas, devices or - # entities within templated services, but searching them should not raise or - # otherwise fail. - assert hass.states.get("script.script_with_templated_services") - for search_type, search_id in ( - ("area", "script.script_with_templated_services"), - ("device", "script.script_with_templated_services"), - ("entity", "script.script_with_templated_services"), - ): - searcher = search.Searcher( - hass, device_registry, entity_registry, entity_sources - ) - assert searcher.async_search(search_type, search_id) == {} - - searcher = search.Searcher(hass, device_registry, entity_registry, entity_sources) - assert searcher.async_search("entity", "light.wled_config_entry_source") == { - "config_entry": {wled_config_entry.entry_id}, - } - - -async def test_area_lookup( - hass: HomeAssistant, - area_registry: ar.AreaRegistry, - device_registry: dr.DeviceRegistry, - entity_registry: er.EntityRegistry, -) -> None: - """Test area based lookup.""" - living_room_area = area_registry.async_create("Living Room") - - await async_setup_component( - hass, - "script", - { - "script": { - "wled": { - "sequence": [ - { - "service": "light.turn_on", - "target": {"area_id": living_room_area.id}, - }, - ] - }, - } + assert search(ItemType.AUTOMATION, "automation.script") == { + ItemType.AREA: {bedroom_area.id, kitchen_area.id, living_room_area.id}, + ItemType.CONFIG_ENTRY: {wled_config_entry.entry_id, hue_config_entry.entry_id}, + ItemType.DEVICE: {wled_device.id, hue_device.id}, + ItemType.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_wled_hue_entity.entity_id, + script_scene_entity.entity_id, }, - ) - - assert await async_setup_component( - hass, - "automation", - { - "automation": [ - { - "alias": "area_turn_on", - "trigger": {"platform": "template", "value_template": "true"}, - "action": [ - { - "service": "light.turn_on", - "data": { - "area_id": living_room_area.id, - }, - }, - ], - }, - ] - }, - ) - - searcher = search.Searcher( - hass, device_registry, entity_registry, MOCK_ENTITY_SOURCES - ) - assert searcher.async_search("area", living_room_area.id) == { - "script": {"script.wled"}, - "automation": {"automation.area_turn_on"}, + ItemType.FLOOR: {first_floor.floor_id, second_floor.floor_id}, + ItemType.SCENE: {scene_wled_hue_entity.entity_id}, + ItemType.SCRIPT: {script_scene_entity.entity_id}, } - searcher = search.Searcher( - hass, device_registry, entity_registry, MOCK_ENTITY_SOURCES - ) - assert searcher.async_search("automation", "automation.area_turn_on") == { - "area": {living_room_area.id}, - } - - -async def test_person_lookup(hass: HomeAssistant) -> None: - """Test searching persons.""" - assert await async_setup_component( - hass, - "person", - { - "person": [ - { - "id": "abcd", - "name": "Paulus", - "device_trackers": ["device_tracker.paulus_iphone"], - } - ] - }, - ) - - device_reg = dr.async_get(hass) - entity_reg = er.async_get(hass) - - searcher = search.Searcher(hass, device_reg, entity_reg, MOCK_ENTITY_SOURCES) - assert searcher.async_search("entity", "device_tracker.paulus_iphone") == { - "person": {"person.paulus"}, - } - - searcher = search.Searcher(hass, device_reg, entity_reg, MOCK_ENTITY_SOURCES) - assert searcher.async_search("entity", "person.paulus") == { - "entity": {"device_tracker.paulus_iphone"}, - } - - -async def test_automation_blueprint(hass): - """Test searching for automation blueprints.""" - - assert await async_setup_component( - hass, - "automation", - { - "automation": [ - { - "alias": "blueprint_automation_1", - "trigger": {"platform": "template", "value_template": "true"}, - "use_blueprint": { - "path": "test_event_service.yaml", - "input": { - "trigger_event": "blueprint_event_1", - "service_to_call": "test.automation_1", - "a_number": 5, - }, - }, - }, - { - "alias": "blueprint_automation_2", - "trigger": {"platform": "template", "value_template": "true"}, - "use_blueprint": { - "path": "test_event_service.yaml", - "input": { - "trigger_event": "blueprint_event_2", - "service_to_call": "test.automation_2", - "a_number": 5, - }, - }, - }, - ] - }, - ) - - # Ensure automations set up correctly. - assert hass.states.get("automation.blueprint_automation_1") is not None - assert hass.states.get("automation.blueprint_automation_1") is not None - - device_reg = dr.async_get(hass) - entity_reg = er.async_get(hass) - - searcher = search.Searcher(hass, device_reg, entity_reg, MOCK_ENTITY_SOURCES) - assert searcher.async_search("automation", "automation.blueprint_automation_1") == { - "automation": {"automation.blueprint_automation_2"}, - "automation_blueprint": {"test_event_service.yaml"}, - "entity": {"light.kitchen"}, - } - - searcher = search.Searcher(hass, device_reg, entity_reg, MOCK_ENTITY_SOURCES) - assert searcher.async_search("automation_blueprint", "test_event_service.yaml") == { - "automation": { + assert not search(ItemType.AUTOMATION_BLUEPRINT, "unknown.yaml") + assert search(ItemType.AUTOMATION_BLUEPRINT, "test_event_service.yaml") == { + ItemType.AUTOMATION: { "automation.blueprint_automation_1", "automation.blueprint_automation_2", + } + } + + assert not search(ItemType.CONFIG_ENTRY, "unknown") + assert search(ItemType.CONFIG_ENTRY, hue_config_entry.entry_id) == { + ItemType.AREA: {kitchen_area.id}, + ItemType.DEVICE: {hue_device.id}, + ItemType.ENTITY: { + hue_segment_1_entity.entity_id, + hue_segment_2_entity.entity_id, }, + ItemType.FLOOR: {first_floor.floor_id}, + ItemType.GROUP: {"group.hue", "group.wled_hue"}, + ItemType.SCENE: {"scene.scene_hue_seg_1", scene_wled_hue_entity.entity_id}, + ItemType.SCRIPT: {"script.device", "script.hue"}, } - - -async def test_script_blueprint(hass): - """Test searching for script blueprints.""" - - assert await async_setup_component( - hass, - "script", - { - "script": { - "blueprint_script_1": { - "use_blueprint": { - "path": "test_service.yaml", - "input": { - "service_to_call": "test.automation", - }, - } - }, - "blueprint_script_2": { - "use_blueprint": { - "path": "test_service.yaml", - "input": { - "service_to_call": "test.automation", - }, - } - }, - } + assert search(ItemType.CONFIG_ENTRY, wled_config_entry.entry_id) == { + ItemType.AREA: {bedroom_area.id, living_room_area.id}, + ItemType.AUTOMATION: {"automation.wled_entity", "automation.wled_device"}, + ItemType.DEVICE: {wled_device.id}, + ItemType.ENTITY: { + wled_segment_1_entity.entity_id, + wled_segment_2_entity.entity_id, }, - ) - - # Ensure automations set up correctly. - assert hass.states.get("script.blueprint_script_1") is not None - assert hass.states.get("script.blueprint_script_1") is not None - - device_reg = dr.async_get(hass) - entity_reg = er.async_get(hass) - - searcher = search.Searcher(hass, device_reg, entity_reg, MOCK_ENTITY_SOURCES) - assert searcher.async_search("script", "script.blueprint_script_1") == { - "entity": {"light.kitchen"}, - "script": {"script.blueprint_script_2"}, - "script_blueprint": {"test_service.yaml"}, + ItemType.FLOOR: {first_floor.floor_id, second_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"}, } - searcher = search.Searcher(hass, device_reg, entity_reg, MOCK_ENTITY_SOURCES) - assert searcher.async_search("script_blueprint", "test_service.yaml") == { - "script": {"script.blueprint_script_1", "script.blueprint_script_2"}, + assert not search(ItemType.DEVICE, "unknown") + assert search(ItemType.DEVICE, wled_device.id) == { + ItemType.AREA: {bedroom_area.id, living_room_area.id}, + ItemType.AUTOMATION: {"automation.wled_entity", "automation.wled_device"}, + ItemType.CONFIG_ENTRY: {wled_config_entry.entry_id}, + ItemType.ENTITY: { + wled_segment_1_entity.entity_id, + wled_segment_2_entity.entity_id, + }, + ItemType.FLOOR: {first_floor.floor_id, second_floor.floor_id}, + ItemType.GROUP: {"group.wled", "group.wled_hue"}, + ItemType.LABEL: {label_christmas.label_id}, + ItemType.SCENE: {"scene.scene_wled_seg_1", scene_wled_hue_entity.entity_id}, + ItemType.SCRIPT: {"script.wled"}, + } + assert search(ItemType.DEVICE, hue_device.id) == { + ItemType.AREA: {kitchen_area.id}, + ItemType.CONFIG_ENTRY: {hue_config_entry.entry_id}, + ItemType.ENTITY: { + hue_segment_1_entity.entity_id, + hue_segment_2_entity.entity_id, + }, + ItemType.FLOOR: {first_floor.floor_id}, + ItemType.GROUP: {"group.hue", "group.wled_hue"}, + ItemType.SCENE: {"scene.scene_hue_seg_1", scene_wled_hue_entity.entity_id}, + ItemType.SCRIPT: {"script.device", "script.hue"}, } + assert not search(ItemType.ENTITY, "sensor.unknown") + assert search(ItemType.ENTITY, wled_segment_1_entity.entity_id) == { + ItemType.AREA: {living_room_area.id}, + ItemType.AUTOMATION: {"automation.wled_entity"}, + ItemType.CONFIG_ENTRY: {wled_config_entry.entry_id}, + ItemType.DEVICE: {wled_device.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.ENTITY, wled_segment_2_entity.entity_id) == { + ItemType.AREA: {bedroom_area.id}, + ItemType.CONFIG_ENTRY: {wled_config_entry.entry_id}, + ItemType.DEVICE: {wled_device.id}, + ItemType.FLOOR: {second_floor.floor_id}, + ItemType.GROUP: {"group.wled", "group.wled_hue"}, + ItemType.SCENE: {scene_wled_hue_entity.entity_id}, + } + assert search(ItemType.ENTITY, hue_segment_1_entity.entity_id) == { + ItemType.AREA: {kitchen_area.id}, + ItemType.CONFIG_ENTRY: {hue_config_entry.entry_id}, + ItemType.DEVICE: {hue_device.id}, + ItemType.FLOOR: {first_floor.floor_id}, + ItemType.GROUP: {"group.hue", "group.wled_hue"}, + ItemType.LABEL: {label_energy.label_id}, + ItemType.SCENE: {"scene.scene_hue_seg_1", scene_wled_hue_entity.entity_id}, + ItemType.SCRIPT: {"script.hue"}, + } + assert search(ItemType.ENTITY, hue_segment_2_entity.entity_id) == { + ItemType.AREA: {kitchen_area.id}, + ItemType.CONFIG_ENTRY: {hue_config_entry.entry_id}, + ItemType.DEVICE: {hue_device.id}, + ItemType.FLOOR: {first_floor.floor_id}, + ItemType.GROUP: {"group.hue", "group.wled_hue"}, + ItemType.SCENE: {scene_wled_hue_entity.entity_id}, + } + assert not search(ItemType.ENTITY, "automation.wled") + assert search(ItemType.ENTITY, script_scene_entity.entity_id) == { + ItemType.AREA: {bedroom_area.id}, + ItemType.AUTOMATION: {"automation.script"}, + ItemType.FLOOR: {second_floor.floor_id}, + ItemType.LABEL: {label_other.label_id}, + ItemType.SCRIPT: {"script.nested"}, + } + assert search(ItemType.ENTITY, "group.wled_hue") == { + ItemType.AUTOMATION: {"automation.group"}, + ItemType.SCRIPT: {"script.group"}, + } + assert search(ItemType.ENTITY, person_paulus_entity.entity_id) == { + ItemType.AREA: {bedroom_area.id}, + ItemType.FLOOR: {second_floor.floor_id}, + ItemType.LABEL: {label_other.label_id}, + } + assert search(ItemType.ENTITY, scene_wled_hue_entity.entity_id) == { + ItemType.AREA: {bedroom_area.id}, + ItemType.AUTOMATION: {"automation.scene"}, + ItemType.FLOOR: {second_floor.floor_id}, + ItemType.LABEL: {label_other.label_id}, + ItemType.SCRIPT: {script_scene_entity.entity_id}, + } + assert search(ItemType.ENTITY, "device_tracker.paulus_iphone") == { + ItemType.PERSON: {person_paulus_entity.entity_id}, + } + assert search(ItemType.ENTITY, "light.wled_config_entry_source") == { + ItemType.CONFIG_ENTRY: {wled_config_entry.entry_id}, + } -async def test_ws_api(hass: HomeAssistant, hass_ws_client: WebSocketGenerator) -> None: - """Test WS API.""" - assert await async_setup_component(hass, "search", {}) + assert not search(ItemType.FLOOR, "unknown") + assert search(ItemType.FLOOR, first_floor.floor_id) == { + ItemType.AREA: {kitchen_area.id, living_room_area.id}, + ItemType.AUTOMATION: { + "automation.area", + "automation.floor", + "automation.wled_device", + }, + ItemType.CONFIG_ENTRY: {hue_config_entry.entry_id, wled_config_entry.entry_id}, + ItemType.DEVICE: {hue_device.id, wled_device.id}, + ItemType.ENTITY: { + wled_segment_1_entity.entity_id, + hue_segment_1_entity.entity_id, + hue_segment_2_entity.entity_id, + }, + ItemType.SCRIPT: {"script.device", "script.area", "script.floor"}, + } + assert search(ItemType.FLOOR, second_floor.floor_id) == { + ItemType.AREA: {bedroom_area.id}, + ItemType.AUTOMATION: {"automation.scene", "automation.script"}, + ItemType.CONFIG_ENTRY: {wled_config_entry.entry_id}, + ItemType.ENTITY: { + wled_segment_2_entity.entity_id, + person_paulus_entity.entity_id, + scene_wled_hue_entity.entity_id, + script_scene_entity.entity_id, + }, + ItemType.GROUP: {"group.wled", "group.wled_hue"}, + ItemType.PERSON: {person_paulus_entity.entity_id}, + ItemType.SCENE: {scene_wled_hue_entity.entity_id}, + ItemType.SCRIPT: {script_scene_entity.entity_id, "script.nested"}, + } - area_reg = ar.async_get(hass) - device_reg = dr.async_get(hass) + assert not search(ItemType.GROUP, "group.unknown") + assert search(ItemType.GROUP, "group.wled") == { + ItemType.AREA: {bedroom_area.id, living_room_area.id}, + ItemType.CONFIG_ENTRY: {wled_config_entry.entry_id}, + ItemType.DEVICE: {wled_device.id}, + ItemType.ENTITY: { + wled_segment_1_entity.entity_id, + wled_segment_2_entity.entity_id, + }, + ItemType.FLOOR: {first_floor.floor_id, second_floor.floor_id}, + } + assert search(ItemType.GROUP, "group.hue") == { + ItemType.AREA: {kitchen_area.id}, + ItemType.CONFIG_ENTRY: {hue_config_entry.entry_id}, + ItemType.DEVICE: {hue_device.id}, + ItemType.ENTITY: { + hue_segment_1_entity.entity_id, + hue_segment_2_entity.entity_id, + }, + ItemType.FLOOR: {first_floor.floor_id}, + } + assert search(ItemType.GROUP, "group.wled_hue") == { + ItemType.AREA: {bedroom_area.id, living_room_area.id, kitchen_area.id}, + ItemType.AUTOMATION: {"automation.group"}, + ItemType.CONFIG_ENTRY: {wled_config_entry.entry_id, hue_config_entry.entry_id}, + ItemType.DEVICE: {wled_device.id, hue_device.id}, + ItemType.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, + }, + ItemType.FLOOR: {first_floor.floor_id, second_floor.floor_id}, + ItemType.SCRIPT: {"script.group"}, + } - kitchen_area = area_reg.async_create("Kitchen") + assert not search(ItemType.LABEL, "unknown") + assert search(ItemType.LABEL, label_christmas.label_id) == { + ItemType.AUTOMATION: {"automation.label"}, + ItemType.DEVICE: {wled_device.id}, + } + assert search(ItemType.LABEL, label_energy.label_id) == { + ItemType.ENTITY: {hue_segment_1_entity.entity_id}, + } + assert search(ItemType.LABEL, label_other.label_id) == { + ItemType.AREA: {bedroom_area.id}, + ItemType.ENTITY: { + scene_wled_hue_entity.entity_id, + person_paulus_entity.entity_id, + script_scene_entity.entity_id, + }, + ItemType.PERSON: {person_paulus_entity.entity_id}, + ItemType.SCENE: {scene_wled_hue_entity.entity_id}, + ItemType.SCRIPT: {"script.label", script_scene_entity.entity_id}, + } - hue_config_entry = MockConfigEntry(domain="hue") - hue_config_entry.add_to_hass(hass) + assert not search(ItemType.PERSON, "person.unknown") + assert search(ItemType.PERSON, person_paulus_entity.entity_id) == { + ItemType.AREA: {bedroom_area.id}, + ItemType.ENTITY: {"device_tracker.paulus_iphone"}, + ItemType.FLOOR: {second_floor.floor_id}, + ItemType.LABEL: {label_other.label_id}, + } - hue_device = device_reg.async_get_or_create( - config_entry_id=hue_config_entry.entry_id, - name="Light Strip", - identifiers=({"hue", "hue-1"}), - ) + assert not search(ItemType.SCENE, "scene.unknown") + assert search(ItemType.SCENE, "scene.scene_wled_seg_1") == { + ItemType.AREA: {living_room_area.id}, + 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}, + } + assert search(ItemType.SCENE, "scene.scene_hue_seg_1") == { + ItemType.AREA: {kitchen_area.id}, + ItemType.CONFIG_ENTRY: {hue_config_entry.entry_id}, + ItemType.DEVICE: {hue_device.id}, + ItemType.ENTITY: {hue_segment_1_entity.entity_id}, + ItemType.FLOOR: {first_floor.floor_id}, + } + assert search(ItemType.SCENE, scene_wled_hue_entity.entity_id) == { + ItemType.AREA: {bedroom_area.id, living_room_area.id, kitchen_area.id}, + ItemType.AUTOMATION: {"automation.scene"}, + ItemType.CONFIG_ENTRY: {wled_config_entry.entry_id, hue_config_entry.entry_id}, + ItemType.DEVICE: {wled_device.id, hue_device.id}, + ItemType.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, + }, + ItemType.FLOOR: {first_floor.floor_id, second_floor.floor_id}, + ItemType.LABEL: {label_other.label_id}, + ItemType.SCRIPT: {script_scene_entity.entity_id}, + } - device_reg.async_update_device(hue_device.id, area_id=kitchen_area.id) + assert not search(ItemType.SCRIPT, "script.unknown") + assert search(ItemType.SCRIPT, "script.blueprint_script_1") == { + ItemType.ENTITY: {"light.kitchen"}, + ItemType.SCRIPT_BLUEPRINT: {"test_service.yaml"}, + } + assert search(ItemType.SCRIPT, "script.blueprint_script_2") == { + ItemType.ENTITY: {"light.kitchen"}, + ItemType.SCRIPT_BLUEPRINT: {"test_service.yaml"}, + } + assert search(ItemType.SCRIPT, "script.wled") == { + ItemType.AREA: {living_room_area.id}, + 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}, + } + assert search(ItemType.SCRIPT, "script.hue") == { + ItemType.AREA: {kitchen_area.id}, + ItemType.CONFIG_ENTRY: {hue_config_entry.entry_id}, + ItemType.DEVICE: {hue_device.id}, + ItemType.ENTITY: {hue_segment_1_entity.entity_id}, + ItemType.FLOOR: {first_floor.floor_id}, + } + assert search(ItemType.SCRIPT, "script.script_with_templated_services") == {} + assert search(ItemType.SCRIPT, "script.device") == { + ItemType.AREA: {kitchen_area.id}, + ItemType.CONFIG_ENTRY: {hue_config_entry.entry_id}, + ItemType.DEVICE: {hue_device.id}, + ItemType.FLOOR: {first_floor.floor_id}, + } + assert search(ItemType.SCRIPT, "script.floor") == { + ItemType.FLOOR: {first_floor.floor_id}, + } + assert search(ItemType.SCRIPT, "script.area") == { + ItemType.AREA: {kitchen_area.id}, + ItemType.FLOOR: {first_floor.floor_id}, + } + assert search(ItemType.SCRIPT, "script.group") == { + ItemType.AREA: {bedroom_area.id, living_room_area.id, kitchen_area.id}, + ItemType.CONFIG_ENTRY: {wled_config_entry.entry_id, hue_config_entry.entry_id}, + ItemType.DEVICE: {wled_device.id, hue_device.id}, + ItemType.ENTITY: { + "group.wled_hue", + wled_segment_1_entity.entity_id, + wled_segment_2_entity.entity_id, + hue_segment_1_entity.entity_id, + hue_segment_2_entity.entity_id, + }, + ItemType.FLOOR: {first_floor.floor_id, second_floor.floor_id}, + ItemType.GROUP: {"group.wled_hue"}, + } + assert search(ItemType.SCRIPT, script_scene_entity.entity_id) == { + ItemType.AREA: {bedroom_area.id, kitchen_area.id, living_room_area.id}, + ItemType.CONFIG_ENTRY: {wled_config_entry.entry_id, hue_config_entry.entry_id}, + ItemType.DEVICE: {wled_device.id, hue_device.id}, + ItemType.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_wled_hue_entity.entity_id, + }, + ItemType.FLOOR: {first_floor.floor_id, second_floor.floor_id}, + ItemType.LABEL: {label_other.label_id}, + ItemType.SCENE: {scene_wled_hue_entity.entity_id}, + } + assert search(ItemType.SCRIPT, "script.nested") == { + ItemType.AREA: {bedroom_area.id, kitchen_area.id, living_room_area.id}, + ItemType.CONFIG_ENTRY: {wled_config_entry.entry_id, hue_config_entry.entry_id}, + ItemType.DEVICE: {wled_device.id, hue_device.id}, + ItemType.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_wled_hue_entity.entity_id, + script_scene_entity.entity_id, + }, + ItemType.FLOOR: {first_floor.floor_id, second_floor.floor_id}, + ItemType.SCENE: {scene_wled_hue_entity.entity_id}, + ItemType.SCRIPT: {script_scene_entity.entity_id}, + } + assert not search(ItemType.SCRIPT_BLUEPRINT, "unknown.yaml") + assert search(ItemType.SCRIPT_BLUEPRINT, "test_service.yaml") == { + ItemType.SCRIPT: {"script.blueprint_script_1", "script.blueprint_script_2"}, + } + + # WebSocket client = await hass_ws_client(hass) - await client.send_json( { "id": 1, @@ -585,6 +948,23 @@ async def test_ws_api(hass: HomeAssistant, hass_ws_client: WebSocketGenerator) - response = await client.receive_json() assert response["success"] assert response["result"] == { - "config_entry": [hue_config_entry.entry_id], - "area": [kitchen_area.id], + ItemType.AREA: [kitchen_area.id], + ItemType.ENTITY: unordered( + [ + hue_segment_1_entity.entity_id, + hue_segment_2_entity.entity_id, + ] + ), + ItemType.GROUP: unordered( + [ + "group.hue", + "group.wled_hue", + ] + ), + ItemType.CONFIG_ENTRY: [hue_config_entry.entry_id], + ItemType.FLOOR: [first_floor.floor_id], + ItemType.SCENE: unordered( + ["scene.scene_hue_seg_1", scene_wled_hue_entity.entity_id] + ), + ItemType.SCRIPT: unordered(["script.device", "script.hue"]), }