From dce20680d7fbf111921b2ce3bba9630422601c0d Mon Sep 17 00:00:00 2001
From: Chris Nussbaum <chrisnussbaum@gmail.com>
Date: Mon, 15 Feb 2021 11:32:22 -0600
Subject: [PATCH] Add duration option to action start deep sleep (#1526)

---
 esphome/components/deep_sleep/__init__.py      | 18 ++++++++++++++----
 .../deep_sleep/deep_sleep_component.h          |  8 +++++++-
 tests/test2.yaml                               |  8 ++++++++
 3 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/esphome/components/deep_sleep/__init__.py b/esphome/components/deep_sleep/__init__.py
index 1b766c9928..9eafcc8609 100644
--- a/esphome/components/deep_sleep/__init__.py
+++ b/esphome/components/deep_sleep/__init__.py
@@ -84,18 +84,28 @@ def to_code(config):
     cg.add_define('USE_DEEP_SLEEP')
 
 
-DEEP_SLEEP_ACTION_SCHEMA = automation.maybe_simple_id({
+DEEP_SLEEP_ENTER_SCHEMA = automation.maybe_simple_id({
+    cv.GenerateID(): cv.use_id(DeepSleepComponent),
+    cv.Optional(CONF_SLEEP_DURATION): cv.positive_time_period_milliseconds,
+})
+
+
+DEEP_SLEEP_PREVENT_SCHEMA = automation.maybe_simple_id({
     cv.GenerateID(): cv.use_id(DeepSleepComponent),
 })
 
 
-@automation.register_action('deep_sleep.enter', EnterDeepSleepAction, DEEP_SLEEP_ACTION_SCHEMA)
+@automation.register_action('deep_sleep.enter', EnterDeepSleepAction, DEEP_SLEEP_ENTER_SCHEMA)
 def deep_sleep_enter_to_code(config, action_id, template_arg, args):
     paren = yield cg.get_variable(config[CONF_ID])
-    yield cg.new_Pvariable(action_id, template_arg, paren)
+    var = cg.new_Pvariable(action_id, template_arg, paren)
+    if CONF_SLEEP_DURATION in config:
+        template_ = yield cg.templatable(config[CONF_SLEEP_DURATION], args, cg.int32)
+        cg.add(var.set_sleep_duration(template_))
+    yield var
 
 
-@automation.register_action('deep_sleep.prevent', PreventDeepSleepAction, DEEP_SLEEP_ACTION_SCHEMA)
+@automation.register_action('deep_sleep.prevent', PreventDeepSleepAction, DEEP_SLEEP_PREVENT_SCHEMA)
 def deep_sleep_prevent_to_code(config, action_id, template_arg, args):
     paren = yield cg.get_variable(config[CONF_ID])
     yield cg.new_Pvariable(action_id, template_arg, paren)
diff --git a/esphome/components/deep_sleep/deep_sleep_component.h b/esphome/components/deep_sleep/deep_sleep_component.h
index 4372a3f66c..09212d7d17 100644
--- a/esphome/components/deep_sleep/deep_sleep_component.h
+++ b/esphome/components/deep_sleep/deep_sleep_component.h
@@ -84,8 +84,14 @@ extern bool global_has_deep_sleep;
 template<typename... Ts> class EnterDeepSleepAction : public Action<Ts...> {
  public:
   EnterDeepSleepAction(DeepSleepComponent *deep_sleep) : deep_sleep_(deep_sleep) {}
+  TEMPLATABLE_VALUE(uint32_t, sleep_duration);
 
-  void play(Ts... x) override { this->deep_sleep_->begin_sleep(true); }
+  void play(Ts... x) override {
+    if (this->sleep_duration_.has_value()) {
+      this->deep_sleep_->set_sleep_duration(this->sleep_duration_.value(x...));
+    }
+    this->deep_sleep_->begin_sleep(true);
+  }
 
  protected:
   DeepSleepComponent *deep_sleep_;
diff --git a/tests/test2.yaml b/tests/test2.yaml
index b109aad758..333025358c 100644
--- a/tests/test2.yaml
+++ b/tests/test2.yaml
@@ -45,6 +45,12 @@ ota:
 logger:
   level: DEBUG
 
+deep_sleep:
+  run_duration: 20s
+  sleep_duration: 50s
+  wakeup_pin: GPIO39
+  wakeup_pin_mode: INVERT_WAKEUP
+
 as3935_i2c:
   irq_pin: GPIO12
 
@@ -302,6 +308,8 @@ text_sensor:
       - homeassistant.tag_scanned:
           tag: 1234-abcd
       - homeassistant.tag_scanned: 1234-abcd
+      - deep_sleep.enter:
+          sleep_duration: 30min
   - platform: template
     name: 'Template Text Sensor'
     lambda: |-