From 3afb564a486b7f4744d58639b073575113b7314b Mon Sep 17 00:00:00 2001
From: Alex <33379584+alexyao2015@users.noreply.github.com>
Date: Tue, 1 Dec 2020 17:41:39 -0500
Subject: [PATCH] Configurable OTA Safe Mode (#1393)

---
 esphome/components/ota/__init__.py     | 9 +++++++--
 esphome/components/ota/ota_component.h | 2 +-
 tests/test1.yaml                       | 2 ++
 tests/test2.yaml                       | 1 +
 tests/test3.yaml                       | 1 +
 5 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/esphome/components/ota/__init__.py b/esphome/components/ota/__init__.py
index e6bcff045f..e4b6946116 100644
--- a/esphome/components/ota/__init__.py
+++ b/esphome/components/ota/__init__.py
@@ -1,6 +1,9 @@
 import esphome.codegen as cg
 import esphome.config_validation as cv
-from esphome.const import CONF_ID, CONF_PASSWORD, CONF_PORT, CONF_SAFE_MODE
+from esphome.const import (
+    CONF_ID, CONF_NUM_ATTEMPTS, CONF_PASSWORD,
+    CONF_PORT, CONF_REBOOT_TIMEOUT, CONF_SAFE_MODE
+)
 from esphome.core import CORE, coroutine_with_priority
 
 CODEOWNERS = ['@esphome/core']
@@ -14,6 +17,8 @@ CONFIG_SCHEMA = cv.Schema({
     cv.Optional(CONF_SAFE_MODE, default=True): cv.boolean,
     cv.SplitDefault(CONF_PORT, esp8266=8266, esp32=3232): cv.port,
     cv.Optional(CONF_PASSWORD, default=''): cv.string,
+    cv.Optional(CONF_REBOOT_TIMEOUT, default='5min'): cv.positive_time_period_milliseconds,
+    cv.Optional(CONF_NUM_ATTEMPTS, default='10'): cv.positive_not_null_int
 }).extend(cv.COMPONENT_SCHEMA)
 
 
@@ -26,7 +31,7 @@ def to_code(config):
     yield cg.register_component(var, config)
 
     if config[CONF_SAFE_MODE]:
-        cg.add(var.start_safe_mode())
+        cg.add(var.start_safe_mode(config[CONF_NUM_ATTEMPTS], config[CONF_REBOOT_TIMEOUT]))
 
     if CORE.is_esp8266:
         cg.add_library('Update', None)
diff --git a/esphome/components/ota/ota_component.h b/esphome/components/ota/ota_component.h
index e37cb7160c..65d44482b8 100644
--- a/esphome/components/ota/ota_component.h
+++ b/esphome/components/ota/ota_component.h
@@ -47,7 +47,7 @@ class OTAComponent : public Component {
   /// Manually set the port OTA should listen on.
   void set_port(uint16_t port);
 
-  void start_safe_mode(uint8_t num_attempts = 10, uint32_t enable_time = 120000);
+  void start_safe_mode(uint8_t num_attempts, uint32_t enable_time);
 
   // ========== INTERNAL METHODS ==========
   // (In most use cases you won't need these)
diff --git a/tests/test1.yaml b/tests/test1.yaml
index 58d7828c1d..ca6ad6fba6 100644
--- a/tests/test1.yaml
+++ b/tests/test1.yaml
@@ -154,6 +154,8 @@ ota:
   safe_mode: True
   password: 'superlongpasswordthatnoonewillknow'
   port: 3286
+  reboot_timeout: 2min
+  num_attempts: 5
 
 logger:
   baud_rate: 0
diff --git a/tests/test2.yaml b/tests/test2.yaml
index a3328efe7b..b975090531 100644
--- a/tests/test2.yaml
+++ b/tests/test2.yaml
@@ -40,6 +40,7 @@ uart:
 ota:
   safe_mode: True
   port: 3286
+  num_attempts: 15
 
 logger:
   level: DEBUG
diff --git a/tests/test3.yaml b/tests/test3.yaml
index 5af006227d..64d5dbfc9b 100644
--- a/tests/test3.yaml
+++ b/tests/test3.yaml
@@ -204,6 +204,7 @@ uart:
 ota:
   safe_mode: True
   port: 3286
+  reboot_timeout: 15min
 
 logger:
   hardware_uart: UART1