diff --git a/homeassistant/helpers/script.py b/homeassistant/helpers/script.py index c43b918c59d..289b649e90c 100644 --- a/homeassistant/helpers/script.py +++ b/homeassistant/helpers/script.py @@ -1063,9 +1063,13 @@ class Script: if self._referenced_areas is not None: return self._referenced_areas - referenced: set[str] = set() + self._referenced_areas: set[str] = set() + Script._find_referenced_areas(self._referenced_areas, self.sequence) + return self._referenced_areas - for step in self.sequence: + @staticmethod + def _find_referenced_areas(referenced, sequence): + for step in sequence: action = cv.determine_script_action(step) if action == cv.SCRIPT_ACTION_CALL_SERVICE: @@ -1076,8 +1080,11 @@ class Script: ): _referenced_extract_ids(data, ATTR_AREA_ID, referenced) - self._referenced_areas = referenced - return referenced + elif action == cv.SCRIPT_ACTION_CHOOSE: + for choice in step[CONF_CHOOSE]: + Script._find_referenced_areas(referenced, choice[CONF_SEQUENCE]) + if CONF_DEFAULT in step: + Script._find_referenced_areas(referenced, step[CONF_DEFAULT]) @property def referenced_devices(self): @@ -1085,9 +1092,13 @@ class Script: if self._referenced_devices is not None: return self._referenced_devices - referenced: set[str] = set() + self._referenced_devices: set[str] = set() + Script._find_referenced_devices(self._referenced_devices, self.sequence) + return self._referenced_devices - for step in self.sequence: + @staticmethod + def _find_referenced_devices(referenced, sequence): + for step in sequence: action = cv.determine_script_action(step) if action == cv.SCRIPT_ACTION_CALL_SERVICE: @@ -1104,8 +1115,13 @@ class Script: elif action == cv.SCRIPT_ACTION_DEVICE_AUTOMATION: referenced.add(step[CONF_DEVICE_ID]) - self._referenced_devices = referenced - return referenced + elif action == cv.SCRIPT_ACTION_CHOOSE: + for choice in step[CONF_CHOOSE]: + for cond in choice[CONF_CONDITIONS]: + referenced |= condition.async_extract_devices(cond) + Script._find_referenced_devices(referenced, choice[CONF_SEQUENCE]) + if CONF_DEFAULT in step: + Script._find_referenced_devices(referenced, step[CONF_DEFAULT]) @property def referenced_entities(self): @@ -1113,9 +1129,13 @@ class Script: if self._referenced_entities is not None: return self._referenced_entities - referenced: set[str] = set() + self._referenced_entities: set[str] = set() + Script._find_referenced_entities(self._referenced_entities, self.sequence) + return self._referenced_entities - for step in self.sequence: + @staticmethod + def _find_referenced_entities(referenced, sequence): + for step in sequence: action = cv.determine_script_action(step) if action == cv.SCRIPT_ACTION_CALL_SERVICE: @@ -1133,8 +1153,14 @@ class Script: elif action == cv.SCRIPT_ACTION_ACTIVATE_SCENE: referenced.add(step[CONF_SCENE]) - self._referenced_entities = referenced - return referenced + elif action == cv.SCRIPT_ACTION_CHOOSE: + for choice in step[CONF_CHOOSE]: + for cond in choice[CONF_CONDITIONS]: + _LOGGER.error("Extracting entities from: %s", cond) + referenced |= condition.async_extract_entities(cond) + Script._find_referenced_entities(referenced, choice[CONF_SEQUENCE]) + if CONF_DEFAULT in step: + Script._find_referenced_entities(referenced, step[CONF_DEFAULT]) def run( self, variables: _VarsType | None = None, context: Context | None = None diff --git a/tests/helpers/test_script.py b/tests/helpers/test_script.py index dfa5ce34ce7..f8d6c0c6e6b 100644 --- a/tests/helpers/test_script.py +++ b/tests/helpers/test_script.py @@ -2282,6 +2282,38 @@ async def test_referenced_entities(hass): }, {"service": "test.script", "data": {"without": "entity_id"}}, {"scene": "scene.hello"}, + { + "choose": [ + { + "conditions": "{{ states.light.choice_1_cond == 'on' }}", + "sequence": [ + { + "service": "test.script", + "data": {"entity_id": "light.choice_1_seq"}, + } + ], + }, + { + "conditions": { + "condition": "state", + "entity_id": "light.choice_2_cond", + "state": "on", + }, + "sequence": [ + { + "service": "test.script", + "data": {"entity_id": "light.choice_2_seq"}, + } + ], + }, + ], + "default": [ + { + "service": "test.script", + "data": {"entity_id": "light.default_seq"}, + } + ], + }, {"event": "test_event"}, {"delay": "{{ delay_period }}"}, ] @@ -2290,13 +2322,18 @@ async def test_referenced_entities(hass): "test_domain", ) assert script_obj.referenced_entities == { - "light.service_not_list", - "light.service_list", - "sensor.condition", - "scene.hello", + # "light.choice_1_cond", # no entity extraction from template conditions + "light.choice_1_seq", + "light.choice_2_cond", + "light.choice_2_seq", + "light.default_seq", "light.direct_entity_referenced", - "light.entity_in_target", "light.entity_in_data_template", + "light.entity_in_target", + "light.service_list", + "light.service_not_list", + "scene.hello", + "sensor.condition", } # Test we cache results. assert script_obj.referenced_entities is script_obj.referenced_entities @@ -2330,19 +2367,60 @@ async def test_referenced_devices(hass): "service": "test.script", "target": {"device_id": ["target-list-id-1", "target-list-id-2"]}, }, + { + "choose": [ + { + "conditions": "{{ is_device_attr('choice-2-cond-dev-id', 'model', 'blah') }}", + "sequence": [ + { + "service": "test.script", + "target": { + "device_id": "choice-1-seq-device-target" + }, + } + ], + }, + { + "conditions": { + "condition": "device", + "device_id": "choice-2-cond-dev-id", + "domain": "switch", + }, + "sequence": [ + { + "service": "test.script", + "target": { + "device_id": "choice-2-seq-device-target" + }, + } + ], + }, + ], + "default": [ + { + "service": "test.script", + "target": {"device_id": "default-device-target"}, + } + ], + }, ] ), "Test Name", "test_domain", ) assert script_obj.referenced_devices == { - "script-dev-id", + # 'choice-1-cond-dev-id', # no device extraction from template conditions + "choice-1-seq-device-target", + "choice-2-cond-dev-id", + "choice-2-seq-device-target", "condition-dev-id", "data-string-id", "data-template-string-id", - "target-string-id", + "default-device-target", + "script-dev-id", "target-list-id-1", "target-list-id-2", + "target-string-id", } # Test we cache results. assert script_obj.referenced_devices is script_obj.referenced_devices