From d05d67414af1202adf62e5175f9f206fdeb818fd Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Tue, 30 May 2023 15:04:35 +0200 Subject: [PATCH] Teach search about blueprints (#78535) --- .../components/automation/__init__.py | 14 +++ homeassistant/components/script/__init__.py | 14 +++ homeassistant/components/search/__init__.py | 33 +++++- tests/components/search/test_init.py | 107 ++++++++++++++++++ 4 files changed, 167 insertions(+), 1 deletion(-) diff --git a/homeassistant/components/automation/__init__.py b/homeassistant/components/automation/__init__.py index 277952c5b2e..600cc6013e4 100644 --- a/homeassistant/components/automation/__init__.py +++ b/homeassistant/components/automation/__init__.py @@ -228,6 +228,20 @@ def automations_with_blueprint(hass: HomeAssistant, blueprint_path: str) -> list ] +@callback +def blueprint_in_automation(hass: HomeAssistant, entity_id: str) -> str | None: + """Return the blueprint the automation is based on or None.""" + if DOMAIN not in hass.data: + return None + + component: EntityComponent[AutomationEntity] = hass.data[DOMAIN] + + if (automation_entity := component.get_entity(entity_id)) is None: + return None + + return automation_entity.referenced_blueprint + + async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: """Set up all automations.""" hass.data[DOMAIN] = component = EntityComponent[AutomationEntity]( diff --git a/homeassistant/components/script/__init__.py b/homeassistant/components/script/__init__.py index 9c4137c1bea..659131e902b 100644 --- a/homeassistant/components/script/__init__.py +++ b/homeassistant/components/script/__init__.py @@ -160,6 +160,20 @@ def scripts_with_blueprint(hass: HomeAssistant, blueprint_path: str) -> list[str ] +@callback +def blueprint_in_script(hass: HomeAssistant, entity_id: str) -> str | None: + """Return the blueprint the script is based on or None.""" + if DOMAIN not in hass.data: + return None + + component: EntityComponent[ScriptEntity] = hass.data[DOMAIN] + + if (script_entity := component.get_entity(entity_id)) is None: + return None + + return script_entity.referenced_blueprint + + async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: """Load the scripts from the configuration.""" hass.data[DOMAIN] = component = EntityComponent[ScriptEntity](LOGGER, DOMAIN, hass) diff --git a/homeassistant/components/search/__init__.py b/homeassistant/components/search/__init__.py index a603d12e8c4..69796800e61 100644 --- a/homeassistant/components/search/__init__.py +++ b/homeassistant/components/search/__init__.py @@ -37,7 +37,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: ( "area", "automation", - "blueprint", + "automation_blueprint", "config_entry", "device", "entity", @@ -45,6 +45,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: "person", "scene", "script", + "script_blueprint", ) ), vol.Required("item_id"): str, @@ -81,10 +82,12 @@ class Searcher: 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"} @@ -176,6 +179,22 @@ class Searcher: for area in automation.areas_in_automation(self.hass, automation_entity_id): self._add_or_resolve("area", area) + if blueprint := automation.blueprint_in_automation( + self.hass, automation_entity_id + ): + self._add_or_resolve("automation_blueprint", blueprint) + + @callback + def _resolve_automation_blueprint(self, blueprint_path) -> 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) + @callback def _resolve_config_entry(self, config_entry_id) -> None: """Resolve a config entry. @@ -295,3 +314,15 @@ class Searcher: for area in script.areas_in_script(self.hass, script_entity_id): self._add_or_resolve("area", area) + + if blueprint := script.blueprint_in_script(self.hass, script_entity_id): + self._add_or_resolve("script_blueprint", blueprint) + + @callback + def _resolve_script_blueprint(self, blueprint_path) -> None: + """Resolve a script blueprint. + + Will only be called if blueprint is an entry point. + """ + for entity_id in script.scripts_with_blueprint(self.hass, blueprint_path): + self._add_or_resolve("script", entity_id) diff --git a/tests/components/search/test_init.py b/tests/components/search/test_init.py index 4ee6f46c2f8..40ec9c22afe 100644 --- a/tests/components/search/test_init.py +++ b/tests/components/search/test_init.py @@ -449,6 +449,113 @@ async def test_person_lookup(hass: HomeAssistant) -> None: } +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": { + "automation.blueprint_automation_1", + "automation.blueprint_automation_2", + }, + } + + +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", + }, + } + }, + } + }, + ) + + # 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"}, + } + + 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"}, + } + + async def test_ws_api(hass: HomeAssistant, hass_ws_client: WebSocketGenerator) -> None: """Test WS API.""" assert await async_setup_component(hass, "search", {})