Skip to content

Commit

Permalink
Teach search about blueprints
Browse files Browse the repository at this point in the history
  • Loading branch information
emontnemery committed Sep 15, 2022
1 parent de7e12e commit 74304a1
Show file tree
Hide file tree
Showing 4 changed files with 167 additions and 1 deletion.
14 changes: 14 additions & 0 deletions homeassistant/components/automation/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -230,6 +230,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](
Expand Down
14 changes: 14 additions & 0 deletions homeassistant/components/script/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -159,6 +159,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)
Expand Down
33 changes: 32 additions & 1 deletion homeassistant/components/search/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,14 +30,15 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
(
"area",
"automation",
"blueprint",
"automation_blueprint",
"config_entry",
"device",
"entity",
"group",
"person",
"scene",
"script",
"script_blueprint",
)
),
vol.Required("item_id"): str,
Expand Down Expand Up @@ -70,10 +71,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"}
Expand Down Expand Up @@ -167,6 +170,22 @@ def _resolve_automation(self, automation_entity_id) -> None:
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.
Expand Down Expand Up @@ -288,3 +307,15 @@ def _resolve_script(self, script_entity_id) -> None:

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)
107 changes: 107 additions & 0 deletions tests/components/search/test_init.py
Original file line number Diff line number Diff line change
Expand Up @@ -398,6 +398,113 @@ async def test_person_lookup(hass):
}


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, hass_ws_client):
"""Test WS API."""
assert await async_setup_component(hass, "search", {})
Expand Down

0 comments on commit 74304a1

Please sign in to comment.