From d52e2019c08943d35994dce43cd91fafdcaad504 Mon Sep 17 00:00:00 2001
From: Flyte <flyte@failcode.co.uk>
Date: Wed, 10 Feb 2016 11:11:02 +0000
Subject: [PATCH] Update mqtt.publish() function to use template_payload.
 Reorganise publish service. Use mqtt.publish() in tests.

---
 homeassistant/components/mqtt/__init__.py | 25 ++++++++++++++---------
 tests/components/test_mqtt.py             | 20 ++++++++----------
 2 files changed, 23 insertions(+), 22 deletions(-)

diff --git a/homeassistant/components/mqtt/__init__.py b/homeassistant/components/mqtt/__init__.py
index 7c390a465d6..5e4e9ba2308 100644
--- a/homeassistant/components/mqtt/__init__.py
+++ b/homeassistant/components/mqtt/__init__.py
@@ -57,11 +57,14 @@ ATTR_RETAIN = 'retain'
 MAX_RECONNECT_WAIT = 300  # seconds
 
 
-def publish(hass, topic, payload, qos=None, retain=None):
+# pylint: disable=too-many-arguments
+def publish(hass, topic, payload=None, qos=None,
+            retain=None, payload_template=None):
     """Publish message to an MQTT topic."""
     data = {
         ATTR_TOPIC: topic,
         ATTR_PAYLOAD: payload,
+        ATTR_PAYLOAD_TEMPLATE: payload_template
     }
     if qos is not None:
         data[ATTR_QOS] = qos
@@ -70,6 +73,7 @@ def publish(hass, topic, payload, qos=None, retain=None):
         data[ATTR_RETAIN] = retain
 
     hass.services.call(DOMAIN, SERVICE_PUBLISH, data)
+# pylint: enable=too-many-arguments
 
 
 def subscribe(hass, topic, callback, qos=DEFAULT_QOS):
@@ -134,24 +138,25 @@ def setup(hass, config):
         """Handle MQTT publish service calls."""
         msg_topic = call.data.get(ATTR_TOPIC)
         payload = call.data.get(ATTR_PAYLOAD)
+        payload_template = call.data.get(ATTR_PAYLOAD_TEMPLATE)
+        qos = call.data.get(ATTR_QOS, DEFAULT_QOS)
+        retain = call.data.get(ATTR_RETAIN, DEFAULT_RETAIN)
         if payload is None:
+            if payload_template is None:
+                _LOGGER.error(
+                    "You must set either '%s' or '%s' to use this service",
+                    ATTR_PAYLOAD, ATTR_PAYLOAD_TEMPLATE)
+                return
             try:
-                payload_template = call.data.get(ATTR_PAYLOAD_TEMPLATE)
-                if payload_template is None:
-                    _LOGGER.error(
-                        "You must set either '%s' or '%s' to use this service",
-                        ATTR_PAYLOAD, ATTR_PAYLOAD_TEMPLATE)
-                    return
                 payload = template.render(hass, payload_template)
-            except AttributeError as exc:
+            except template.jinja2.TemplateError as exc:
                 _LOGGER.error(
                     "Unable to publish to '%s': rendering payload template of "
                     "'%s' failed because %s.",
                     msg_topic, payload_template, exc)
+                return
         if msg_topic is None or payload is None:
             return
-        qos = call.data.get(ATTR_QOS, DEFAULT_QOS)
-        retain = call.data.get(ATTR_RETAIN, DEFAULT_RETAIN)
         MQTT_CLIENT.publish(msg_topic, payload, qos, retain)
 
     hass.bus.listen_once(EVENT_HOMEASSISTANT_START, start_mqtt)
diff --git a/tests/components/test_mqtt.py b/tests/components/test_mqtt.py
index 5746dd5a273..de8192a6b10 100644
--- a/tests/components/test_mqtt.py
+++ b/tests/components/test_mqtt.py
@@ -80,10 +80,9 @@ class TestMQTT(unittest.TestCase):
         """
         If 'payload_template' is provided and 'payload' is not, then render it.
         """
-        self.hass.services.call(mqtt.DOMAIN, mqtt.SERVICE_PUBLISH, {
-            mqtt.ATTR_TOPIC: "test/topic",
-            mqtt.ATTR_PAYLOAD_TEMPLATE: "{{ 1+1 }}"
-        }, blocking=True)
+        mqtt.publish(self.hass, "test/topic",
+                     **{mqtt.ATTR_PAYLOAD_TEMPLATE: "{{ 1+1 }}"})
+        self.hass.pool.block_till_done()
         self.assertTrue(mqtt.MQTT_CLIENT.publish.called)
         self.assertEqual(mqtt.MQTT_CLIENT.publish.call_args[0][1], "2")
 
@@ -93,11 +92,9 @@ class TestMQTT(unittest.TestCase):
         """
         payload = "not a template"
         payload_template = "a template"
-        self.hass.services.call(mqtt.DOMAIN, mqtt.SERVICE_PUBLISH, {
-            mqtt.ATTR_TOPIC: "test/topic",
-            mqtt.ATTR_PAYLOAD: payload,
-            mqtt.ATTR_PAYLOAD_TEMPLATE: payload_template
-        }, blocking=True)
+        mqtt.publish(self.hass, "test/topic", payload,
+                     **{mqtt.ATTR_PAYLOAD_TEMPLATE: payload_template})
+        self.hass.pool.block_till_done()
         self.assertTrue(mqtt.MQTT_CLIENT.publish.called)
         self.assertEqual(mqtt.MQTT_CLIENT.publish.call_args[0][1], payload)
 
@@ -105,9 +102,8 @@ class TestMQTT(unittest.TestCase):
         """
         If neither 'payload' or 'payload_template' is provided then fail.
         """
-        self.hass.services.call(mqtt.DOMAIN, mqtt.SERVICE_PUBLISH, {
-            mqtt.ATTR_TOPIC: "test/topic"
-        }, blocking=True)
+        mqtt.publish(self.hass, "test/topic")
+        self.hass.pool.block_till_done()
         self.assertFalse(mqtt.MQTT_CLIENT.publish.called)
 
     def test_subscribe_topic(self):