From 67680bcfa8f436e1cace14a4b695da0ccabff3d6 Mon Sep 17 00:00:00 2001 From: Paulus Schoutsen Date: Wed, 5 Feb 2020 07:52:21 -0800 Subject: [PATCH] Automation device/entity extraction to include triggers + conditions (#31474) * Add support for extracting triggers * Add support for extracting triggers * Fix test --- .../components/automation/__init__.py | 172 ++++++++++++------ tests/components/automation/test_init.py | 24 ++- tests/components/search/test_init.py | 4 +- 3 files changed, 143 insertions(+), 57 deletions(-) diff --git a/homeassistant/components/automation/__init__.py b/homeassistant/components/automation/__init__.py index 45f892d783e..528a314dd7b 100644 --- a/homeassistant/components/automation/__init__.py +++ b/homeassistant/components/automation/__init__.py @@ -1,16 +1,19 @@ """Allow to set up simple automation rules via the config file.""" -from functools import partial import importlib import logging -from typing import Any, Awaitable, Callable, List +from typing import Any, Awaitable, Callable, List, Optional, Set import voluptuous as vol +from homeassistant.components import sun from homeassistant.const import ( ATTR_ENTITY_ID, ATTR_NAME, + CONF_DEVICE_ID, + CONF_ENTITY_ID, CONF_ID, CONF_PLATFORM, + CONF_ZONE, EVENT_AUTOMATION_TRIGGERED, EVENT_HOMEASSISTANT_START, SERVICE_RELOAD, @@ -130,7 +133,7 @@ def automations_with_entity(hass: HomeAssistant, entity_id: str) -> List[str]: results = [] for automation_entity in component.entities: - if entity_id in automation_entity.action_script.referenced_entities: + if entity_id in automation_entity.referenced_entities: results.append(automation_entity.entity_id) return results @@ -149,7 +152,7 @@ def entities_in_automation(hass: HomeAssistant, entity_id: str) -> List[str]: if automation_entity is None: return [] - return list(automation_entity.action_script.referenced_entities) + return list(automation_entity.referenced_entities) @callback @@ -163,7 +166,7 @@ def automations_with_device(hass: HomeAssistant, device_id: str) -> List[str]: results = [] for automation_entity in component.entities: - if device_id in automation_entity.action_script.referenced_devices: + if device_id in automation_entity.referenced_devices: results.append(automation_entity.entity_id) return results @@ -182,7 +185,7 @@ def devices_in_automation(hass: HomeAssistant, entity_id: str) -> List[str]: if automation_entity is None: return [] - return list(automation_entity.action_script.referenced_devices) + return list(automation_entity.referenced_devices) async def async_setup(hass, config): @@ -232,7 +235,7 @@ class AutomationEntity(ToggleEntity, RestoreEntity): self, automation_id, name, - async_attach_triggers, + trigger_config, cond_func, action_script, hidden, @@ -241,7 +244,7 @@ class AutomationEntity(ToggleEntity, RestoreEntity): """Initialize an automation entity.""" self._id = automation_id self._name = name - self._async_attach_triggers = async_attach_triggers + self._trigger_config = trigger_config self._async_detach_triggers = None self._cond_func = cond_func self.action_script = action_script @@ -249,6 +252,8 @@ class AutomationEntity(ToggleEntity, RestoreEntity): self._hidden = hidden self._initial_state = initial_state self._is_enabled = False + self._referenced_entities: Optional[Set[str]] = None + self._referenced_devices: Optional[Set[str]] = None @property def name(self): @@ -280,6 +285,45 @@ class AutomationEntity(ToggleEntity, RestoreEntity): """Return True if entity is on.""" return self._async_detach_triggers is not None or self._is_enabled + @property + def referenced_devices(self): + """Return a set of referenced devices.""" + if self._referenced_devices is not None: + return self._referenced_devices + + referenced = self.action_script.referenced_devices + + if self._cond_func is not None: + for conf in self._cond_func.config: + referenced |= condition.async_extract_devices(conf) + + for conf in self._trigger_config: + device = _trigger_extract_device(conf) + if device is not None: + referenced.add(device) + + self._referenced_devices = referenced + return referenced + + @property + def referenced_entities(self): + """Return a set of referenced entities.""" + if self._referenced_entities is not None: + return self._referenced_entities + + referenced = self.action_script.referenced_entities + + if self._cond_func is not None: + for conf in self._cond_func.config: + referenced |= condition.async_extract_entities(conf) + + for conf in self._trigger_config: + for entity_id in _trigger_extract_entities(conf): + referenced.add(entity_id) + + self._referenced_entities = referenced + return referenced + async def async_added_to_hass(self) -> None: """Startup with initial state or previous state.""" await super().async_added_to_hass() @@ -330,7 +374,11 @@ class AutomationEntity(ToggleEntity, RestoreEntity): This method is a coroutine. """ - if not skip_condition and not self._cond_func(variables): + if ( + not skip_condition + and self._cond_func is not None + and not self._cond_func(variables) + ): return # Create a new context referring to the old context. @@ -373,9 +421,7 @@ class AutomationEntity(ToggleEntity, RestoreEntity): # HomeAssistant is starting up if self.hass.state != CoreState.not_running: - self._async_detach_triggers = await self._async_attach_triggers( - self.async_trigger - ) + self._async_detach_triggers = await self._async_attach_triggers() self.async_write_ha_state() return @@ -385,9 +431,7 @@ class AutomationEntity(ToggleEntity, RestoreEntity): if not self._is_enabled or self._async_detach_triggers is not None: return - self._async_detach_triggers = await self._async_attach_triggers( - self.async_trigger - ) + self._async_detach_triggers = await self._async_attach_triggers() self.hass.bus.async_listen_once( EVENT_HOMEASSISTANT_START, async_enable_automation @@ -407,6 +451,38 @@ class AutomationEntity(ToggleEntity, RestoreEntity): self.async_write_ha_state() + async def _async_attach_triggers(self): + """Set up the triggers.""" + removes = [] + info = {"name": self._name} + + for conf in self._trigger_config: + platform = importlib.import_module( + ".{}".format(conf[CONF_PLATFORM]), __name__ + ) + + remove = await platform.async_attach_trigger( + self.hass, conf, self.async_trigger, info + ) + + if not remove: + _LOGGER.error("Error setting up trigger %s", self._name) + continue + + _LOGGER.info("Initialized trigger %s", self._name) + removes.append(remove) + + if not removes: + return None + + @callback + def remove_triggers(): + """Remove attached triggers.""" + for remove in removes: + remove() + + return remove_triggers + @property def device_state_attributes(self): """Return automation attributes.""" @@ -441,22 +517,12 @@ async def _async_process_config(hass, config, component): if cond_func is None: continue else: + cond_func = None - def cond_func(variables): - """Condition will always pass.""" - return True - - async_attach_triggers = partial( - _async_process_trigger, - hass, - config, - config_block.get(CONF_TRIGGER, []), - name, - ) entity = AutomationEntity( automation_id, name, - async_attach_triggers, + config_block[CONF_TRIGGER], cond_func, action_script, hidden, @@ -471,7 +537,7 @@ async def _async_process_config(hass, config, component): async def _async_process_if(hass, config, p_config): """Process if checks.""" - if_configs = p_config.get(CONF_CONDITION) + if_configs = p_config[CONF_CONDITION] checks = [] for if_config in if_configs: @@ -485,35 +551,33 @@ async def _async_process_if(hass, config, p_config): """AND all conditions.""" return all(check(hass, variables) for check in checks) + if_action.config = if_configs + return if_action -async def _async_process_trigger(hass, config, trigger_configs, name, action): - """Set up the triggers. - - This method is a coroutine. - """ - removes = [] - info = {"name": name} - - for conf in trigger_configs: - platform = importlib.import_module(".{}".format(conf[CONF_PLATFORM]), __name__) - - remove = await platform.async_attach_trigger(hass, conf, action, info) - - if not remove: - _LOGGER.error("Error setting up trigger %s", name) - continue - - _LOGGER.info("Initialized trigger %s", name) - removes.append(remove) - - if not removes: +@callback +def _trigger_extract_device(trigger_conf: dict) -> Optional[str]: + """Extract devices from a trigger config.""" + if trigger_conf[CONF_PLATFORM] != "device": return None - def remove_triggers(): - """Remove attached triggers.""" - for remove in removes: - remove() + return trigger_conf[CONF_DEVICE_ID] - return remove_triggers + +@callback +def _trigger_extract_entities(trigger_conf: dict) -> List[str]: + """Extract entities from a trigger config.""" + if trigger_conf[CONF_PLATFORM] in ("state", "numeric_state"): + return trigger_conf[CONF_ENTITY_ID] + + if trigger_conf[CONF_PLATFORM] == "zone": + return trigger_conf[CONF_ENTITY_ID] + [trigger_conf[CONF_ZONE]] + + if trigger_conf[CONF_PLATFORM] == "geo_location": + return [trigger_conf[CONF_ZONE]] + + if trigger_conf[CONF_PLATFORM] == "sun": + return [sun.ENTITY_ID] + + return [] diff --git a/tests/components/automation/test_init.py b/tests/components/automation/test_init.py index 391c9646dd4..c27a0262a4e 100644 --- a/tests/components/automation/test_init.py +++ b/tests/components/automation/test_init.py @@ -935,6 +935,11 @@ async def test_extraction_functions(hass): { "alias": "test1", "trigger": {"platform": "state", "entity_id": "sensor.trigger_1"}, + "condition": { + "condition": "state", + "entity_id": "light.condition_state", + "state": "on", + }, "action": [ { "service": "test.script", @@ -954,7 +959,20 @@ async def test_extraction_functions(hass): }, { "alias": "test2", - "trigger": {"platform": "state", "entity_id": "sensor.trigger_2"}, + "trigger": { + "platform": "device", + "domain": "light", + "type": "turned_on", + "entity_id": "light.trigger_2", + "device_id": "trigger-device-2", + }, + "condition": { + "condition": "device", + "device_id": "condition-device", + "domain": "light", + "type": "is_on", + "entity_id": "light.bla", + }, "action": [ { "service": "test.script", @@ -989,6 +1007,8 @@ async def test_extraction_functions(hass): "automation.test2", } assert set(automation.entities_in_automation(hass, "automation.test1")) == { + "sensor.trigger_1", + "light.condition_state", "light.in_both", "light.in_first", } @@ -997,6 +1017,8 @@ async def test_extraction_functions(hass): "automation.test2", } assert set(automation.devices_in_automation(hass, "automation.test2")) == { + "trigger-device-2", + "condition-device", "device-in-both", "device-in-last", } diff --git a/tests/components/search/test_init.py b/tests/components/search/test_init.py index 54a32bed229..a379b91f82a 100644 --- a/tests/components/search/test_init.py +++ b/tests/components/search/test_init.py @@ -163,7 +163,7 @@ async def test_search(hass): "automation": [ { "alias": "wled_entity", - "trigger": {"platform": "state", "entity_id": "sensor.trigger_1"}, + "trigger": {"platform": "template", "value_template": "true"}, "action": [ { "service": "test.script", @@ -173,7 +173,7 @@ async def test_search(hass): }, { "alias": "wled_device", - "trigger": {"platform": "state", "entity_id": "sensor.trigger_1"}, + "trigger": {"platform": "template", "value_template": "true"}, "action": [ { "domain": "light",