Fix extracting entity and device IDs from scripts (#44048)

* Fix extracting entity and device IDs from scripts

* Fix extracting from data_template

homeassistant/helpers/script.py

@@ -22,10 +22,10 @@ from async_timeout import timeout
 import voluptuous as vol
 
 from homeassistant import exceptions
-import homeassistant.components.device_automation as device_automation
+from homeassistant.components import device_automation, scene
 from homeassistant.components.logger import LOGSEVERITY
-import homeassistant.components.scene as scene
 from homeassistant.const import (
+    ATTR_DEVICE_ID,
     ATTR_ENTITY_ID,
     CONF_ALIAS,
     CONF_CHOOSE,
@@ -44,6 +44,7 @@ from homeassistant.const import (
     CONF_REPEAT,
     CONF_SCENE,
     CONF_SEQUENCE,
+    CONF_TARGET,
     CONF_TIMEOUT,
     CONF_UNTIL,
     CONF_VARIABLES,
@@ -60,13 +61,9 @@ from homeassistant.core import (
     HomeAssistant,
     callback,
 )
-from homeassistant.helpers import condition, config_validation as cv, template
+from homeassistant.helpers import condition, config_validation as cv, service, template
 from homeassistant.helpers.event import async_call_later, async_track_template
 from homeassistant.helpers.script_variables import ScriptVariables
-from homeassistant.helpers.service import (
-    CONF_SERVICE_DATA,
-    async_prepare_call_from_config,
-)
 from homeassistant.helpers.trigger import (
     async_initialize_triggers,
     async_validate_trigger_config,
@@ -429,13 +426,13 @@ class _ScriptRun:
         self._script.last_action = self._action.get(CONF_ALIAS, "call service")
         self._log("Executing step %s", self._script.last_action)
 
-        domain, service, service_data = async_prepare_call_from_config(
+        domain, service_name, service_data = service.async_prepare_call_from_config(
             self._hass, self._action, self._variables
         )
 
         running_script = (
             domain == "automation"
-            and service == "trigger"
+            and service_name == "trigger"
             or domain in ("python_script", "script")
         )
         # If this might start a script then disable the call timeout.
@@ -448,7 +445,7 @@ class _ScriptRun:
         service_task = self._hass.async_create_task(
             self._hass.services.async_call(
                 domain,
-                service,
+                service_name,
                 service_data,
                 blocking=True,
                 context=self._context,
@@ -755,6 +752,23 @@ async def _async_stop_scripts_at_shutdown(hass, event):
 _VarsType = Union[Dict[str, Any], MappingProxyType]
 
 
+def _referenced_extract_ids(data: Dict, key: str, found: Set[str]) -> None:
+    """Extract referenced IDs."""
+    if not data:
+        return
+
+    item_ids = data.get(key)
+
+    if item_ids is None or isinstance(item_ids, template.Template):
+        return
+
+    if isinstance(item_ids, str):
+        item_ids = [item_ids]
+
+    for item_id in item_ids:
+        found.add(item_id)
+
+
 class Script:
     """Representation of a script."""
 
@@ -889,7 +903,16 @@ class Script:
         for step in self.sequence:
             action = cv.determine_script_action(step)
 
-            if action == cv.SCRIPT_ACTION_CHECK_CONDITION:
+            if action == cv.SCRIPT_ACTION_CALL_SERVICE:
+                for data in (
+                    step,
+                    step.get(CONF_TARGET),
+                    step.get(service.CONF_SERVICE_DATA),
+                    step.get(service.CONF_SERVICE_DATA_TEMPLATE),
+                ):
+                    _referenced_extract_ids(data, ATTR_DEVICE_ID, referenced)
+
+            elif action == cv.SCRIPT_ACTION_CHECK_CONDITION:
                 referenced |= condition.async_extract_devices(step)
 
             elif action == cv.SCRIPT_ACTION_DEVICE_AUTOMATION:
@@ -910,20 +933,13 @@ class Script:
             action = cv.determine_script_action(step)
 
             if action == cv.SCRIPT_ACTION_CALL_SERVICE:
-                data = step.get(CONF_SERVICE_DATA)
-                if not data:
-                    continue
-
-                entity_ids = data.get(ATTR_ENTITY_ID)
-
-                if entity_ids is None or isinstance(entity_ids, template.Template):
-                    continue
-
-                if isinstance(entity_ids, str):
-                    entity_ids = [entity_ids]
-
-                for entity_id in entity_ids:
-                    referenced.add(entity_id)
+                for data in (
+                    step,
+                    step.get(CONF_TARGET),
+                    step.get(service.CONF_SERVICE_DATA),
+                    step.get(service.CONF_SERVICE_DATA_TEMPLATE),
+                ):
+                    _referenced_extract_ids(data, ATTR_ENTITY_ID, referenced)
 
             elif action == cv.SCRIPT_ACTION_CHECK_CONDITION:
                 referenced |= condition.async_extract_entities(step)