From d249820bcdd6483e0ab170f4ac0536495a49cb55 Mon Sep 17 00:00:00 2001
From: Otto Winter <otto@otto-winter.com>
Date: Wed, 6 Jun 2018 08:35:23 +0200
Subject: [PATCH] Fix MQTT availability option

---
 esphomeyaml/config_validation.py | 2 +-
 esphomeyaml/helpers.py           | 7 +++++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/esphomeyaml/config_validation.py b/esphomeyaml/config_validation.py
index 6ae3932271..24f13a8fbd 100644
--- a/esphomeyaml/config_validation.py
+++ b/esphomeyaml/config_validation.py
@@ -547,7 +547,7 @@ MQTT_COMPONENT_SCHEMA = vol.Schema({
     vol.Optional(CONF_RETAIN): boolean,
     vol.Optional(CONF_DISCOVERY): boolean,
     vol.Optional(CONF_STATE_TOPIC): publish_topic,
-    vol.Optional(CONF_AVAILABILITY): MQTT_COMPONENT_AVAILABILITY_SCHEMA,
+    vol.Optional(CONF_AVAILABILITY): vol.Any(None, MQTT_COMPONENT_AVAILABILITY_SCHEMA),
 })
 
 MQTT_COMMAND_COMPONENT_SCHEMA = MQTT_COMPONENT_SCHEMA.extend({
diff --git a/esphomeyaml/helpers.py b/esphomeyaml/helpers.py
index 9c5467bedd..5da6e6d406 100644
--- a/esphomeyaml/helpers.py
+++ b/esphomeyaml/helpers.py
@@ -604,8 +604,11 @@ def setup_mqtt_component(obj, config):
         add(obj.set_custom_command_topic(config[CONF_COMMAND_TOPIC]))
     if CONF_AVAILABILITY in config:
         availability = config[CONF_AVAILABILITY]
-        add(obj.set_availability(availability[CONF_TOPIC], availability[CONF_PAYLOAD_AVAILABLE],
-                                 availability[CONF_PAYLOAD_NOT_AVAILABLE]))
+        if not availability:
+            add(obj.disable_availability())
+        else:
+            add(obj.set_availability(availability[CONF_TOPIC], availability[CONF_PAYLOAD_AVAILABLE],
+                                     availability[CONF_PAYLOAD_NOT_AVAILABLE]))
 
 
 # shlex's quote for Python 2.7