Support target triggers in automation relation extraction (#160369)

Author: abmantis

homeassistant/components/automation/__init__.py

@@ -7,7 +7,7 @@
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass
 import logging
-from typing import Any, Protocol, cast
+from typing import Any, Literal, Protocol, cast
 
 from propcache.api import cached_property
 import voluptuous as vol
@@ -16,7 +16,10 @@
 from homeassistant.components.blueprint import CONF_USE_BLUEPRINT
 from homeassistant.components.labs import async_listen as async_labs_listen
 from homeassistant.const import (
+    ATTR_AREA_ID,
     ATTR_ENTITY_ID,
+    ATTR_FLOOR_ID,
+    ATTR_LABEL_ID,
     ATTR_MODE,
     ATTR_NAME,
     CONF_ACTIONS,
@@ -30,6 +33,7 @@
     CONF_OPTIONS,
     CONF_PATH,
     CONF_PLATFORM,
+    CONF_TARGET,
     CONF_TRIGGERS,
     CONF_VARIABLES,
     CONF_ZONE,
@@ -588,20 +592,32 @@ def is_on(self) -> bool:
         """Return True if entity is on."""
         return self._async_detach_triggers is not None or self._is_enabled
 
-    @property
+    @cached_property
     def referenced_labels(self) -> set[str]:
         """Return a set of referenced labels."""
-        return self.action_script.referenced_labels
+        referenced = self.action_script.referenced_labels
 
-    @property
+        for conf in self._trigger_config:
+            referenced |= set(_get_targets_from_trigger_config(conf, ATTR_LABEL_ID))
+        return referenced
+
+    @cached_property
     def referenced_floors(self) -> set[str]:
         """Return a set of referenced floors."""
-        return self.action_script.referenced_floors
+        referenced = self.action_script.referenced_floors
+
+        for conf in self._trigger_config:
+            referenced |= set(_get_targets_from_trigger_config(conf, ATTR_FLOOR_ID))
+        return referenced
 
     @cached_property
     def referenced_areas(self) -> set[str]:
         """Return a set of referenced areas."""
-        return self.action_script.referenced_areas
+        referenced = self.action_script.referenced_areas
+
+        for conf in self._trigger_config:
+            referenced |= set(_get_targets_from_trigger_config(conf, ATTR_AREA_ID))
+        return referenced
 
     @property
     def referenced_blueprint(self) -> str | None:
@@ -1209,6 +1225,9 @@ def _trigger_extract_devices(trigger_conf: dict) -> list[str]:
     if trigger_conf[CONF_PLATFORM] == "tag" and CONF_DEVICE_ID in trigger_conf:
         return trigger_conf[CONF_DEVICE_ID]  # type: ignore[no-any-return]
 
+    if target_devices := _get_targets_from_trigger_config(trigger_conf, CONF_DEVICE_ID):
+        return target_devices
+
     return []
 
 
@@ -1239,9 +1258,28 @@ def _trigger_extract_entities(trigger_conf: dict) -> list[str]:
     ):
         return [trigger_conf[CONF_EVENT_DATA][CONF_ENTITY_ID]]
 
+    if target_entities := _get_targets_from_trigger_config(
+        trigger_conf, CONF_ENTITY_ID
+    ):
+        return target_entities
+
     return []
 
 
+@callback
+def _get_targets_from_trigger_config(
+    config: dict,
+    target: Literal["entity_id", "device_id", "area_id", "floor_id", "label_id"],
+) -> list[str]:
+    """Extract targets from a target config."""
+    if not (target_conf := config.get(CONF_TARGET)):
+        return []
+    if not (targets := target_conf.get(target)):
+        return []
+
+    return [targets] if isinstance(targets, str) else targets
+
+
 @websocket_api.websocket_command({"type": "automation/config", "entity_id": str})
 def websocket_config(
     hass: HomeAssistant,

tests/components/automation/test_init.py

@@ -2232,6 +2232,202 @@ async def test_extraction_functions(
     assert automation.blueprint_in_automation(hass, "automation.test3") is None
 
 
+async def test_extraction_functions_with_targets(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    hass_ws_client: WebSocketGenerator,
+) -> None:
+    """Test extraction functions with targets in triggers.
+
+    This test verifies that targets specified in trigger configurations
+    (using new-style triggers that support target) are properly extracted for
+    entity, device, area, floor, and label references.
+    """
+    config_entry = MockConfigEntry(domain="fake_integration", data={})
+    config_entry.mock_state(hass, ConfigEntryState.LOADED)
+    config_entry.add_to_hass(hass)
+
+    trigger_device = device_registry.async_get_or_create(
+        config_entry_id=config_entry.entry_id,
+        connections={(dr.CONNECTION_NETWORK_MAC, "00:00:00:00:00:01")},
+    )
+
+    await async_setup_component(hass, "homeassistant", {})
+    await async_setup_component(
+        hass, "scene", {"scene": {"name": "test", "entities": {}}}
+    )
+    await hass.async_block_till_done()
+
+    # Enable the new_triggers_conditions feature flag to allow new-style triggers
+    assert await async_setup_component(hass, "labs", {})
+    ws_client = await hass_ws_client(hass)
+    await ws_client.send_json_auto_id(
+        {
+            "type": "labs/update",
+            "domain": "automation",
+            "preview_feature": "new_triggers_conditions",
+            "enabled": True,
+        }
+    )
+    msg = await ws_client.receive_json()
+    assert msg["success"]
+    await hass.async_block_till_done()
+
+    assert await async_setup_component(
+        hass,
+        DOMAIN,
+        {
+            DOMAIN: [
+                {
+                    "alias": "test1",
+                    "triggers": [
+                        # Single entity_id in target
+                        {
+                            "trigger": "scene.activated",
+                            "target": {"entity_id": "scene.target_entity"},
+                        },
+                        # Multiple entity_ids in target
+                        {
+                            "trigger": "scene.activated",
+                            "target": {
+                                "entity_id": [
+                                    "scene.target_entity_list1",
+                                    "scene.target_entity_list2",
+                                ]
+                            },
+                        },
+                        # Single device_id in target
+                        {
+                            "trigger": "scene.activated",
+                            "target": {"device_id": trigger_device.id},
+                        },
+                        # Multiple device_ids in target
+                        {
+                            "trigger": "scene.activated",
+                            "target": {
+                                "device_id": [
+                                    "target-device-1",
+                                    "target-device-2",
+                                ]
+                            },
+                        },
+                        # Single area_id in target
+                        {
+                            "trigger": "scene.activated",
+                            "target": {"area_id": "area-target-single"},
+                        },
+                        # Multiple area_ids in target
+                        {
+                            "trigger": "scene.activated",
+                            "target": {"area_id": ["area-target-1", "area-target-2"]},
+                        },
+                        # Single floor_id in target
+                        {
+                            "trigger": "scene.activated",
+                            "target": {"floor_id": "floor-target-single"},
+                        },
+                        # Multiple floor_ids in target
+                        {
+                            "trigger": "scene.activated",
+                            "target": {
+                                "floor_id": ["floor-target-1", "floor-target-2"]
+                            },
+                        },
+                        # Single label_id in target
+                        {
+                            "trigger": "scene.activated",
+                            "target": {"label_id": "label-target-single"},
+                        },
+                        # Multiple label_ids in target
+                        {
+                            "trigger": "scene.activated",
+                            "target": {
+                                "label_id": ["label-target-1", "label-target-2"]
+                            },
+                        },
+                        # Combined targets
+                        {
+                            "trigger": "scene.activated",
+                            "target": {
+                                "entity_id": "scene.combined_entity",
+                                "device_id": "combined-device",
+                                "area_id": "combined-area",
+                                "floor_id": "combined-floor",
+                                "label_id": "combined-label",
+                            },
+                        },
+                    ],
+                    "conditions": [],
+                    "actions": [
+                        {
+                            "action": "test.script",
+                            "data": {"entity_id": "light.action_entity"},
+                        },
+                    ],
+                },
+            ]
+        },
+    )
+
+    # Test entity extraction from trigger targets
+    assert set(automation.entities_in_automation(hass, "automation.test1")) == {
+        "scene.target_entity",
+        "scene.target_entity_list1",
+        "scene.target_entity_list2",
+        "scene.combined_entity",
+        "light.action_entity",
+    }
+
+    # Test device extraction from trigger targets
+    assert set(automation.devices_in_automation(hass, "automation.test1")) == {
+        trigger_device.id,
+        "target-device-1",
+        "target-device-2",
+        "combined-device",
+    }
+
+    # Test area extraction from trigger targets
+    assert set(automation.areas_in_automation(hass, "automation.test1")) == {
+        "area-target-single",
+        "area-target-1",
+        "area-target-2",
+        "combined-area",
+    }
+
+    # Test floor extraction from trigger targets
+    assert set(automation.floors_in_automation(hass, "automation.test1")) == {
+        "floor-target-single",
+        "floor-target-1",
+        "floor-target-2",
+        "combined-floor",
+    }
+
+    # Test label extraction from trigger targets
+    assert set(automation.labels_in_automation(hass, "automation.test1")) == {
+        "label-target-single",
+        "label-target-1",
+        "label-target-2",
+        "combined-label",
+    }
+
+    # Test automations_with_* functions
+    assert set(automation.automations_with_entity(hass, "scene.target_entity")) == {
+        "automation.test1"
+    }
+    assert set(automation.automations_with_device(hass, trigger_device.id)) == {
+        "automation.test1"
+    }
+    assert set(automation.automations_with_area(hass, "area-target-single")) == {
+        "automation.test1"
+    }
+    assert set(automation.automations_with_floor(hass, "floor-target-single")) == {
+        "automation.test1"
+    }
+    assert set(automation.automations_with_label(hass, "label-target-single")) == {
+        "automation.test1"
+    }
+
+
 async def test_logbook_humanify_automation_triggered_event(hass: HomeAssistant) -> None:
     """Test humanifying Automation Trigger event."""
     hass.config.components.add("recorder")