From 3ec931ffa49a2e93681af6fb416535030efc4b3c Mon Sep 17 00:00:00 2001
From: Otto Winter <otto@otto-winter.com>
Date: Fri, 9 Nov 2018 20:05:50 +0100
Subject: [PATCH] Add restore state option to template switch (#222)

* Add restore state option to template switch

* Add test
---
 esphomeyaml/components/switch/template.py | 8 +++++++-
 esphomeyaml/const.py                      | 1 +
 tests/test1.yaml                          | 2 ++
 3 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/esphomeyaml/components/switch/template.py b/esphomeyaml/components/switch/template.py
index 0d7f967322..121bbeff1c 100644
--- a/esphomeyaml/components/switch/template.py
+++ b/esphomeyaml/components/switch/template.py
@@ -4,18 +4,21 @@ import esphomeyaml.config_validation as cv
 from esphomeyaml import automation
 from esphomeyaml.components import switch
 from esphomeyaml.const import CONF_LAMBDA, CONF_MAKE_ID, CONF_NAME, CONF_TURN_OFF_ACTION, \
-    CONF_TURN_ON_ACTION, CONF_OPTIMISTIC
+    CONF_TURN_ON_ACTION, CONF_OPTIMISTIC, CONF_RESTORE_STATE
 from esphomeyaml.helpers import App, Application, process_lambda, variable, NoArg, add, bool_, \
     optional
 
 MakeTemplateSwitch = Application.MakeTemplateSwitch
+TemplateSwitch = switch.switch_ns.TemplateSwitch
 
 PLATFORM_SCHEMA = cv.nameable(switch.SWITCH_PLATFORM_SCHEMA.extend({
+    cv.GenerateID(): cv.declare_variable_id(TemplateSwitch),
     cv.GenerateID(CONF_MAKE_ID): cv.declare_variable_id(MakeTemplateSwitch),
     vol.Optional(CONF_LAMBDA): cv.lambda_,
     vol.Optional(CONF_OPTIMISTIC): cv.boolean,
     vol.Optional(CONF_TURN_OFF_ACTION): automation.validate_automation(single=True),
     vol.Optional(CONF_TURN_ON_ACTION): automation.validate_automation(single=True),
+    vol.Optional(CONF_RESTORE_STATE): cv.boolean,
 }), cv.has_at_least_one_key(CONF_LAMBDA, CONF_OPTIMISTIC))
 
 
@@ -40,6 +43,9 @@ def to_code(config):
     if CONF_OPTIMISTIC in config:
         add(make.Ptemplate_.set_optimistic(config[CONF_OPTIMISTIC]))
 
+    if CONF_RESTORE_STATE in config:
+        add(make.Ptemplate_.set_restore_state(config[CONF_RESTORE_STATE]))
+
 
 BUILD_FLAGS = '-DUSE_TEMPLATE_SWITCH'
 
diff --git a/esphomeyaml/const.py b/esphomeyaml/const.py
index e289bd10a2..0eacce7392 100644
--- a/esphomeyaml/const.py
+++ b/esphomeyaml/const.py
@@ -359,6 +359,7 @@ CONF_STEP_PIN = 'step_pin'
 CONF_DIR_PIN = 'dir_pin'
 CONF_SLEEP_PIN = 'sleep_pin'
 CONF_SEND_FIRST_AT = 'send_first_at'
+CONF_RESTORE_STATE = 'restore_state'
 
 ALLOWED_NAME_CHARS = u'abcdefghijklmnopqrstuvwxyz0123456789_'
 ARDUINO_VERSION_ESP32_DEV = 'https://github.com/platformio/platform-espressif32.git#feature/stage'
diff --git a/tests/test1.yaml b/tests/test1.yaml
index 9d7ca7cb8a..ee8073aa82 100644
--- a/tests/test1.yaml
+++ b/tests/test1.yaml
@@ -806,6 +806,7 @@ switch:
         level: !lambda 'return 0.5;'
     turn_off_action:
     - switch.turn_on: living_room_lights_off
+    restore_state: False
   - platform: restart
     name: "Living Room Restart"
   - platform: shutdown
@@ -832,6 +833,7 @@ switch:
         // Switch is OFF, do something else here
       }
     optimistic: true
+    restore_state: True
   - platform: uart
     name: "UART String Output"
     data: 'DataToSend'