Automation device/entity extraction to include triggers + conditions (#31474)

* Add support for extracting triggers

* Add support for extracting triggers

* Fix test

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_detach_triggers = await self._async_attach_triggers()
-                self.async_trigger
-            )
             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_detach_triggers = await self._async_attach_triggers()
-                self.async_trigger
-            )
 
         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):
+@callback
-    """Set up the triggers.
+def _trigger_extract_device(trigger_conf: dict) -> Optional[str]:
-
+    """Extract devices from a trigger config."""
-    This method is a coroutine.
+    if trigger_conf[CONF_PLATFORM] != "device":
-    """
-    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:
         return None
 
-    def remove_triggers():
+    return trigger_conf[CONF_DEVICE_ID]
-        """Remove attached triggers."""
-        for remove in removes:
-            remove()
 
-    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 []

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",
     }

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",