WaitSensor - time_to_wait templated

providers/standard/src/airflow/providers/standard/sensors/time_delta.py

@@ -18,6 +18,7 @@
 from __future__ import annotations
 
 import warnings
+from collections.abc import Sequence
 from datetime import datetime, timedelta
 from time import sleep
 from typing import TYPE_CHECKING, Any
@@ -177,6 +178,8 @@ class WaitSensor(BaseSensorOperator):
     :param deferrable: Run sensor in deferrable mode
     """
 
+    template_fields: Sequence[str] = ("time_to_wait",)
+
     def __init__(
         self,
         time_to_wait: timedelta | int,
@@ -185,20 +188,24 @@ def __init__(
     ) -> None:
         super().__init__(**kwargs)
         self.deferrable = deferrable
-        if isinstance(time_to_wait, int):
-            self.time_to_wait = timedelta(minutes=time_to_wait)
-        else:
-            self.time_to_wait = time_to_wait
+        self.time_to_wait = time_to_wait
+
+    def _resolve_time_to_wait(self) -> timedelta:
+        value = self.time_to_wait
+        if isinstance(value, timedelta):
+            return value
+        return timedelta(minutes=int(value))
 
     def execute(self, context: Context) -> None:
+        time_to_wait = self._resolve_time_to_wait()
         if self.deferrable:
             self.defer(
                 trigger=(
-                    TimeDeltaTrigger(self.time_to_wait, end_from_trigger=True)
+                    TimeDeltaTrigger(time_to_wait, end_from_trigger=True)
                     if AIRFLOW_V_3_0_PLUS
-                    else TimeDeltaTrigger(self.time_to_wait)
+                    else TimeDeltaTrigger(time_to_wait)
                 ),
                 method_name="execute_complete",
             )
         else:
-            sleep(int(self.time_to_wait.total_seconds()))
+            sleep(int(time_to_wait.total_seconds()))

providers/standard/tests/unit/standard/sensors/test_time_delta.py

@@ -17,6 +17,7 @@
 # under the License.
 from __future__ import annotations
 
+import re
 from datetime import timedelta
 from typing import Any
 
@@ -253,3 +254,27 @@ def test_timedelta_sensor_async_run_after_vs_interval(self, run_after, interval_
             op.execute(context)
 
         assert caught.value.trigger.moment == expected_time
+
+    @pytest.mark.parametrize(
+        "time_to_wait",
+        [timedelta(minutes=1), 1, "{{ 1*2 }}"],
+    )
+    def test_wait_sensor_templating(self, mocker, time_to_wait):
+        defer_mock = mocker.patch(DEFER_PATH)
+        op = WaitSensor(task_id="wait_sensor_check", time_to_wait=time_to_wait, dag=self.dag, deferrable=True)
+
+        with time_machine.travel(pendulum.datetime(year=2024, month=8, day=1, tz="UTC"), tick=False):
+            context = op.render_template_fields({})
+            op.execute(context)
+            defer_mock.assert_called_once()
+
+    def test_wait_sensor_templating_error(self, mocker):
+        op = WaitSensor(
+            task_id="wait_sensor_check", time_to_wait="{{ 'nothing' }}", dag=self.dag, deferrable=True
+        )
+        with time_machine.travel(pendulum.datetime(year=2024, month=8, day=1, tz="UTC"), tick=False):
+            context = op.render_template_fields({})
+            with pytest.raises(
+                ValueError, match=re.escape("invalid literal for int() with base 10: 'nothing'")
+            ):
+                op.execute(context)