diff --git a/esphomeyaml/components/wifi.py b/esphomeyaml/components/wifi.py
index 4a030a6a94..0a95bd6040 100644
--- a/esphomeyaml/components/wifi.py
+++ b/esphomeyaml/components/wifi.py
@@ -135,7 +135,7 @@ def to_code(config):
         add(wifi.set_reboot_timeout(config[CONF_REBOOT_TIMEOUT]))
 
     if CONF_POWER_SAVE_MODE in config:
-        add(wifi.set_power_save_mode(WIFI_POWER_SAVE_MODES[CONF_POWER_SAVE_MODE]))
+        add(wifi.set_power_save_mode(WIFI_POWER_SAVE_MODES[config[CONF_POWER_SAVE_MODE]]))
 
 
 def lib_deps(config):
diff --git a/tests/test1.yaml b/tests/test1.yaml
index 997de4aa48..d6135207cc 100644
--- a/tests/test1.yaml
+++ b/tests/test1.yaml
@@ -35,6 +35,7 @@ wifi:
   hostname: myverylonghostname
   domain: .local
   reboot_timeout: 120s
+  power_save_mode: none
 
 mqtt:
   broker: '192.168.178.84'