From f59e727f162b73a0707f3a0bbc1cdb0cb84bad5a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Thomas=20Lov=C3=A9n?= Date: Fri, 11 Sep 2020 13:16:25 +0200 Subject: [PATCH] Set variable values in scripts (#39915) Co-authored-by: Paulus Schoutsen --- homeassistant/helpers/config_validation.py | 13 ++++++ homeassistant/helpers/script.py | 9 ++++ homeassistant/helpers/script_variables.py | 25 +++++++---- tests/helpers/test_script.py | 39 ++++++++++++++++ tests/helpers/test_script_variables.py | 52 ++++++++++++++++++++++ 5 files changed, 129 insertions(+), 9 deletions(-) diff --git a/homeassistant/helpers/config_validation.py b/homeassistant/helpers/config_validation.py index 602a8ebfd2a..282e63e6440 100644 --- a/homeassistant/helpers/config_validation.py +++ b/homeassistant/helpers/config_validation.py @@ -67,6 +67,7 @@ from homeassistant.const import ( CONF_UNIT_SYSTEM_METRIC, CONF_UNTIL, CONF_VALUE_TEMPLATE, + CONF_VARIABLES, CONF_WAIT_FOR_TRIGGER, CONF_WAIT_TEMPLATE, CONF_WHILE, @@ -1127,6 +1128,13 @@ _SCRIPT_WAIT_FOR_TRIGGER_SCHEMA = vol.Schema( } ) +_SCRIPT_SET_SCHEMA = vol.Schema( + { + vol.Optional(CONF_ALIAS): string, + vol.Required(CONF_VARIABLES): SCRIPT_VARIABLES_SCHEMA, + } +) + SCRIPT_ACTION_DELAY = "delay" SCRIPT_ACTION_WAIT_TEMPLATE = "wait_template" SCRIPT_ACTION_CHECK_CONDITION = "condition" @@ -1137,6 +1145,7 @@ SCRIPT_ACTION_ACTIVATE_SCENE = "scene" SCRIPT_ACTION_REPEAT = "repeat" SCRIPT_ACTION_CHOOSE = "choose" SCRIPT_ACTION_WAIT_FOR_TRIGGER = "wait_for_trigger" +SCRIPT_ACTION_VARIABLES = "variables" def determine_script_action(action: dict) -> str: @@ -1168,6 +1177,9 @@ def determine_script_action(action: dict) -> str: if CONF_WAIT_FOR_TRIGGER in action: return SCRIPT_ACTION_WAIT_FOR_TRIGGER + if CONF_VARIABLES in action: + return SCRIPT_ACTION_VARIABLES + return SCRIPT_ACTION_CALL_SERVICE @@ -1182,4 +1194,5 @@ ACTION_TYPE_SCHEMAS: Dict[str, Callable[[Any], dict]] = { SCRIPT_ACTION_REPEAT: _SCRIPT_REPEAT_SCHEMA, SCRIPT_ACTION_CHOOSE: _SCRIPT_CHOOSE_SCHEMA, SCRIPT_ACTION_WAIT_FOR_TRIGGER: _SCRIPT_WAIT_FOR_TRIGGER_SCHEMA, + SCRIPT_ACTION_VARIABLES: _SCRIPT_SET_SCHEMA, } diff --git a/homeassistant/helpers/script.py b/homeassistant/helpers/script.py index bd1442587eb..717e9c3980c 100644 --- a/homeassistant/helpers/script.py +++ b/homeassistant/helpers/script.py @@ -46,6 +46,7 @@ from homeassistant.const import ( CONF_SEQUENCE, CONF_TIMEOUT, CONF_UNTIL, + CONF_VARIABLES, CONF_WAIT_FOR_TRIGGER, CONF_WAIT_TEMPLATE, CONF_WHILE, @@ -612,6 +613,14 @@ class _ScriptRun: task.cancel() remove_triggers() + async def _async_variables_step(self): + """Set a variable value.""" + self._script.last_action = self._action.get(CONF_ALIAS, "setting variables") + self._log("Executing step %s", self._script.last_action) + self._variables = self._action[CONF_VARIABLES].async_render( + self._hass, self._variables, render_as_defaults=False + ) + async def _async_run_script(self, script): """Execute a script.""" await self._async_run_long_action( diff --git a/homeassistant/helpers/script_variables.py b/homeassistant/helpers/script_variables.py index 001c3b8667c..3140fc4dced 100644 --- a/homeassistant/helpers/script_variables.py +++ b/homeassistant/helpers/script_variables.py @@ -19,21 +19,31 @@ class ScriptVariables: self, hass: HomeAssistant, run_variables: Optional[Mapping[str, Any]], + *, + render_as_defaults: bool = True, ) -> Dict[str, Any]: """Render script variables. - The run variables are used to compute the static variables, but afterwards will also - be merged on top of the static variables. + The run variables are used to compute the static variables. + + If `render_as_defaults` is True, the run variables will not be overridden. + """ if self._has_template is None: self._has_template = template.is_complex(self.variables) template.attach(hass, self.variables) if not self._has_template: - rendered_variables = dict(self.variables) + if render_as_defaults: + rendered_variables = dict(self.variables) - if run_variables is not None: - rendered_variables.update(run_variables) + if run_variables is not None: + rendered_variables.update(run_variables) + else: + rendered_variables = ( + {} if run_variables is None else dict(run_variables) + ) + rendered_variables.update(self.variables) return rendered_variables @@ -42,14 +52,11 @@ class ScriptVariables: for key, value in self.variables.items(): # We can skip if we're going to override this key with # run variables anyway - if key in rendered_variables: + if render_as_defaults and key in rendered_variables: continue rendered_variables[key] = template.render_complex(value, rendered_variables) - if run_variables: - rendered_variables.update(run_variables) - return rendered_variables def as_dict(self) -> dict: diff --git a/tests/helpers/test_script.py b/tests/helpers/test_script.py index d298283d11e..0bd353e1fa0 100644 --- a/tests/helpers/test_script.py +++ b/tests/helpers/test_script.py @@ -1785,3 +1785,42 @@ async def test_started_action(hass, caplog): await hass.async_block_till_done() assert log_message in caplog.text + + +async def test_set_variable(hass, caplog): + """Test setting variables in scripts.""" + sequence = cv.SCRIPT_SCHEMA( + [ + {"variables": {"variable": "value"}}, + {"service": "test.script", "data": {"value": "{{ variable }}"}}, + ] + ) + script_obj = script.Script(hass, sequence, "test script", "test_domain") + + mock_calls = async_mock_service(hass, "test", "script") + + await script_obj.async_run(context=Context()) + await hass.async_block_till_done() + + assert mock_calls[0].data["value"] == "value" + + +async def test_set_redefines_variable(hass, caplog): + """Test setting variables based on their current value.""" + sequence = cv.SCRIPT_SCHEMA( + [ + {"variables": {"variable": "1"}}, + {"service": "test.script", "data": {"value": "{{ variable }}"}}, + {"variables": {"variable": "{{ variable | int + 1 }}"}}, + {"service": "test.script", "data": {"value": "{{ variable }}"}}, + ] + ) + script_obj = script.Script(hass, sequence, "test script", "test_domain") + + mock_calls = async_mock_service(hass, "test", "script") + + await script_obj.async_run(context=Context()) + await hass.async_block_till_done() + + assert mock_calls[0].data["value"] == "1" + assert mock_calls[1].data["value"] == "2" diff --git a/tests/helpers/test_script_variables.py b/tests/helpers/test_script_variables.py index 6e671d14a23..20a70cb33eb 100644 --- a/tests/helpers/test_script_variables.py +++ b/tests/helpers/test_script_variables.py @@ -24,6 +24,28 @@ async def test_static_vars_run_args(): assert orig == orig_copy +async def test_static_vars_no_default(): + """Test static vars.""" + orig = {"hello": "world"} + var = cv.SCRIPT_VARIABLES_SCHEMA(orig) + rendered = var.async_render(None, None, render_as_defaults=False) + assert rendered is not orig + assert rendered == orig + + +async def test_static_vars_run_args_no_default(): + """Test static vars.""" + orig = {"hello": "world"} + orig_copy = dict(orig) + var = cv.SCRIPT_VARIABLES_SCHEMA(orig) + rendered = var.async_render( + None, {"hello": "override", "run": "var"}, render_as_defaults=False + ) + assert rendered == {"hello": "world", "run": "var"} + # Make sure we don't change original vars + assert orig == orig_copy + + async def test_template_vars(hass): """Test template vars.""" var = cv.SCRIPT_VARIABLES_SCHEMA({"hello": "{{ 1 + 1 }}"}) @@ -53,6 +75,36 @@ async def test_template_vars_run_args(hass): } +async def test_template_vars_no_default(hass): + """Test template vars.""" + var = cv.SCRIPT_VARIABLES_SCHEMA({"hello": "{{ 1 + 1 }}"}) + rendered = var.async_render(hass, None, render_as_defaults=False) + assert rendered == {"hello": "2"} + + +async def test_template_vars_run_args_no_default(hass): + """Test template vars.""" + var = cv.SCRIPT_VARIABLES_SCHEMA( + { + "something": "{{ run_var_ex + 1 }}", + "something_2": "{{ run_var_ex + 1 }}", + } + ) + rendered = var.async_render( + hass, + { + "run_var_ex": 5, + "something_2": 1, + }, + render_as_defaults=False, + ) + assert rendered == { + "run_var_ex": 5, + "something": "6", + "something_2": "6", + } + + async def test_template_vars_error(hass): """Test template vars.""" var = cv.SCRIPT_VARIABLES_SCHEMA({"hello": "{{ canont.work }}"})