From 2ef25e7414fdff80f0e600e41620aa04ab3babe8 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Fri, 1 Jan 2021 02:03:34 -1000
Subject: [PATCH] Fix script wait templates with now/utcnow (#44717)

---
 homeassistant/helpers/event.py  |  6 +++---
 homeassistant/helpers/script.py | 15 +++++++++++----
 tests/helpers/test_event.py     | 27 +++++++++++++++++++++++++++
 tests/helpers/test_script.py    | 26 ++++++++++++++++++++++++++
 4 files changed, 67 insertions(+), 7 deletions(-)

diff --git a/homeassistant/helpers/event.py b/homeassistant/helpers/event.py
index 661e1a11b56..f06ac8aca3f 100644
--- a/homeassistant/helpers/event.py
+++ b/homeassistant/helpers/event.py
@@ -715,9 +715,9 @@ def async_track_template(
 
         hass.async_run_hass_job(
             job,
-            event.data.get("entity_id"),
-            event.data.get("old_state"),
-            event.data.get("new_state"),
+            event and event.data.get("entity_id"),
+            event and event.data.get("old_state"),
+            event and event.data.get("new_state"),
         )
 
     info = async_track_template_result(
diff --git a/homeassistant/helpers/script.py b/homeassistant/helpers/script.py
index 48a662e3a81..77c842a27fe 100644
--- a/homeassistant/helpers/script.py
+++ b/homeassistant/helpers/script.py
@@ -62,7 +62,11 @@ from homeassistant.core import (
     callback,
 )
 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.event import (
+    TrackTemplate,
+    async_call_later,
+    async_track_template_result,
+)
 from homeassistant.helpers.script_variables import ScriptVariables
 from homeassistant.helpers.trigger import (
     async_initialize_triggers,
@@ -355,7 +359,7 @@ class _ScriptRun:
             return
 
         @callback
-        def async_script_wait(entity_id, from_s, to_s):
+        def _async_script_wait(event, updates):
             """Handle script after template condition is true."""
             self._variables["wait"] = {
                 "remaining": to_context.remaining if to_context else delay,
@@ -364,9 +368,12 @@ class _ScriptRun:
             done.set()
 
         to_context = None
-        unsub = async_track_template(
-            self._hass, wait_template, async_script_wait, self._variables
+        info = async_track_template_result(
+            self._hass,
+            [TrackTemplate(wait_template, self._variables)],
+            _async_script_wait,
         )
+        unsub = info.async_remove
 
         self._changed()
         done = asyncio.Event()
diff --git a/tests/helpers/test_event.py b/tests/helpers/test_event.py
index 1b9ddc52191..9c7ddb09f85 100644
--- a/tests/helpers/test_event.py
+++ b/tests/helpers/test_event.py
@@ -908,6 +908,33 @@ async def test_track_template_error_can_recover(hass, caplog):
     assert "UndefinedError" not in caplog.text
 
 
+async def test_track_template_time_change(hass, caplog):
+    """Test tracking template with time change."""
+    template_error = Template("{{ utcnow().minute % 2 == 0 }}", hass)
+    calls = []
+
+    @ha.callback
+    def error_callback(entity_id, old_state, new_state):
+        calls.append((entity_id, old_state, new_state))
+
+    start_time = dt_util.utcnow() + timedelta(hours=24)
+    time_that_will_not_match_right_away = start_time.replace(minute=1, second=0)
+    with patch(
+        "homeassistant.util.dt.utcnow", return_value=time_that_will_not_match_right_away
+    ):
+        async_track_template(hass, template_error, error_callback)
+        await hass.async_block_till_done()
+        assert not calls
+
+    first_time = start_time.replace(minute=2, second=0)
+    with patch("homeassistant.util.dt.utcnow", return_value=first_time):
+        async_fire_time_changed(hass, first_time)
+        await hass.async_block_till_done()
+
+    assert len(calls) == 1
+    assert calls[0] == (None, None, None)
+
+
 async def test_track_template_result(hass):
     """Test tracking template."""
     specific_runs = []
diff --git a/tests/helpers/test_script.py b/tests/helpers/test_script.py
index 92666335f28..c81ed681d42 100644
--- a/tests/helpers/test_script.py
+++ b/tests/helpers/test_script.py
@@ -780,6 +780,32 @@ async def test_wait_template_variables_in(hass):
         assert not script_obj.is_running
 
 
+async def test_wait_template_with_utcnow(hass):
+    """Test the wait template with utcnow."""
+    sequence = cv.SCRIPT_SCHEMA({"wait_template": "{{ utcnow().hours == 12 }}"})
+    script_obj = script.Script(hass, sequence, "Test Name", "test_domain")
+    wait_started_flag = async_watch_for_action(script_obj, "wait")
+    start_time = dt_util.utcnow() + timedelta(hours=24)
+
+    try:
+        hass.async_create_task(script_obj.async_run(context=Context()))
+        async_fire_time_changed(hass, start_time.replace(hour=5))
+        assert not script_obj.is_running
+        async_fire_time_changed(hass, start_time.replace(hour=12))
+
+        await asyncio.wait_for(wait_started_flag.wait(), 1)
+
+        assert script_obj.is_running
+    except (AssertionError, asyncio.TimeoutError):
+        await script_obj.async_stop()
+        raise
+    else:
+        async_fire_time_changed(hass, start_time.replace(hour=3))
+        await hass.async_block_till_done()
+
+        assert not script_obj.is_running
+
+
 @pytest.mark.parametrize("mode", ["no_timeout", "timeout_finish", "timeout_not_finish"])
 @pytest.mark.parametrize("action_type", ["template", "trigger"])
 async def test_wait_variables_out(hass, mode, action_type):