From 34db9d9ef2b0748abee992958e783ec88a69398b Mon Sep 17 00:00:00 2001
From: Jesse Hills <3060199+jesserockz@users.noreply.github.com>
Date: Wed, 13 Oct 2021 08:23:24 +1300
Subject: [PATCH] Add optional timeout for wait_until action (#2282)

---
 esphome/automation.py          |  7 +++++++
 esphome/core/base_automation.h | 12 ++++++++++++
 tests/test1.yaml               |  5 ++++-
 tests/test3.yaml               | 12 +++++++-----
 4 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/esphome/automation.py b/esphome/automation.py
index 71c564b906..0768bf8869 100644
--- a/esphome/automation.py
+++ b/esphome/automation.py
@@ -6,6 +6,7 @@ from esphome.const import (
     CONF_ELSE,
     CONF_ID,
     CONF_THEN,
+    CONF_TIMEOUT,
     CONF_TRIGGER_ID,
     CONF_TYPE_ID,
     CONF_TIME,
@@ -244,6 +245,9 @@ def validate_wait_until(value):
     schema = cv.Schema(
         {
             cv.Required(CONF_CONDITION): validate_potentially_and_condition,
+            cv.Optional(CONF_TIMEOUT): cv.templatable(
+                cv.positive_time_period_milliseconds
+            ),
         }
     )
     if isinstance(value, dict) and CONF_CONDITION in value:
@@ -255,6 +259,9 @@ def validate_wait_until(value):
 async def wait_until_action_to_code(config, action_id, template_arg, args):
     conditions = await build_condition(config[CONF_CONDITION], template_arg, args)
     var = cg.new_Pvariable(action_id, template_arg, conditions)
+    if CONF_TIMEOUT in config:
+        template_ = await cg.templatable(config[CONF_TIMEOUT], args, cg.uint32)
+        cg.add(var.set_timeout_value(template_))
     await cg.register_component(var, {})
     return var
 
diff --git a/esphome/core/base_automation.h b/esphome/core/base_automation.h
index fa49786d1d..d97d369d33 100644
--- a/esphome/core/base_automation.h
+++ b/esphome/core/base_automation.h
@@ -228,6 +228,8 @@ template<typename... Ts> class WaitUntilAction : public Action<Ts...>, public Co
  public:
   WaitUntilAction(Condition<Ts...> *condition) : condition_(condition) {}
 
+  TEMPLATABLE_VALUE(uint32_t, timeout_value)
+
   void play_complex(Ts... x) override {
     this->num_running_++;
     // Check if we can continue immediately.
@@ -238,6 +240,12 @@ template<typename... Ts> class WaitUntilAction : public Action<Ts...>, public Co
       return;
     }
     this->var_ = std::make_tuple(x...);
+
+    if (this->timeout_value_.has_value()) {
+      auto f = std::bind(&WaitUntilAction<Ts...>::play_next_, this, x...);
+      this->set_timeout("timeout", this->timeout_value_.value(x...), f);
+    }
+
     this->loop();
   }
 
@@ -249,6 +257,8 @@ template<typename... Ts> class WaitUntilAction : public Action<Ts...>, public Co
       return;
     }
 
+    this->cancel_timeout("timeout");
+
     this->play_next_tuple_(this->var_);
   }
 
@@ -257,6 +267,8 @@ template<typename... Ts> class WaitUntilAction : public Action<Ts...>, public Co
   void play(Ts... x) override { /* ignore - see play_complex */
   }
 
+  void stop() override { this->cancel_timeout("timeout"); }
+
  protected:
   Condition<Ts...> *condition_;
   std::tuple<Ts...> var_{};
diff --git a/tests/test1.yaml b/tests/test1.yaml
index fd142a63fd..540a715dde 100644
--- a/tests/test1.yaml
+++ b/tests/test1.yaml
@@ -364,8 +364,11 @@ sensor:
       then:
         - lambda: >-
             ESP_LOGD("main", "Got value range %f", x);
+        - wait_until: wifi.connected
         - wait_until:
-            binary_sensor.is_on: binary_sensor1
+            condition:
+              binary_sensor.is_on: binary_sensor1
+            timeout: 1s
     on_raw_value:
       - lambda: >-
           ESP_LOGD("main", "Got raw value %f", x);
diff --git a/tests/test3.yaml b/tests/test3.yaml
index 9d2839a026..73e314c94c 100644
--- a/tests/test3.yaml
+++ b/tests/test3.yaml
@@ -5,10 +5,13 @@ esphome:
   board: d1_mini
   build_path: build/test3
   on_boot:
-    - wait_until:
-        - api.connected
-        - wifi.connected
-        - time.has_time
+    - if:
+        condition:
+          - api.connected
+          - wifi.connected
+          - time.has_time
+        then:
+          - logger.log: "Have time"
   includes:
     - custom.h
 
@@ -1291,4 +1294,3 @@ dsmr:
 daly_bms:
   update_interval: 20s
   uart_id: uart1
-