Merge pull request #1138 from flyte/mqtt-publish-template

Add template support to mqtt.publish service payload.

homeassistant/components/mqtt/__init__.py

@@ -14,6 +14,7 @@ import time
 
 from homeassistant.exceptions import HomeAssistantError
 import homeassistant.util as util
+from homeassistant.util import template
 from homeassistant.helpers import validate_config
 from homeassistant.const import (
     EVENT_HOMEASSISTANT_START, EVENT_HOMEASSISTANT_STOP)
@@ -49,24 +50,34 @@ DEFAULT_PROTOCOL = PROTOCOL_311
 
 ATTR_TOPIC = 'topic'
 ATTR_PAYLOAD = 'payload'
+ATTR_PAYLOAD_TEMPLATE = 'payload_template'
 ATTR_QOS = 'qos'
 ATTR_RETAIN = 'retain'
 
 MAX_RECONNECT_WAIT = 300  # seconds
 
 
-def publish(hass, topic, payload, qos=None, retain=None):
-    """Publish message to an MQTT topic."""
-    data = {
-        ATTR_TOPIC: topic,
-        ATTR_PAYLOAD: payload,
-    }
+def _build_publish_data(topic, qos, retain):
+    """Build the arguments for the publish service without the payload."""
+    data = {ATTR_TOPIC: topic}
     if qos is not None:
         data[ATTR_QOS] = qos
-
     if retain is not None:
         data[ATTR_RETAIN] = retain
+    return data
 
+
+def publish(hass, topic, payload, qos=None, retain=None):
+    """Publish message to an MQTT topic."""
+    data = _build_publish_data(topic, qos, retain)
+    data[ATTR_PAYLOAD] = payload
+    hass.services.call(DOMAIN, SERVICE_PUBLISH, data)
+
+
+def publish_template(hass, topic, payload_template, qos=None, retain=None):
+    """Publish message to an MQTT topic using a template payload."""
+    data = _build_publish_data(topic, qos, retain)
+    data[ATTR_PAYLOAD_TEMPLATE] = payload_template
     hass.services.call(DOMAIN, SERVICE_PUBLISH, data)
 
 
@@ -132,8 +143,23 @@ 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.render(hass, payload_template)
+            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
         MQTT_CLIENT.publish(msg_topic, payload, qos, retain)

tests/components/test_mqtt.py

@@ -68,7 +68,7 @@ class TestMQTT(unittest.TestCase):
         self.assertEqual('test-payload',
                          self.calls[0][0].data['service_data'][mqtt.ATTR_PAYLOAD])
 
-    def test_service_call_without_topic_does_not_publush(self):
+    def test_service_call_without_topic_does_not_publish(self):
         self.hass.bus.fire(EVENT_CALL_SERVICE, {
             ATTR_DOMAIN: mqtt.DOMAIN,
             ATTR_SERVICE: mqtt.SERVICE_PUBLISH
@@ -76,6 +76,42 @@ class TestMQTT(unittest.TestCase):
         self.hass.pool.block_till_done()
         self.assertTrue(not mqtt.MQTT_CLIENT.publish.called)
 
+    def test_service_call_with_template_payload_renders_template(self):
+        """
+        If 'payload_template' is provided and 'payload' is not, then render it.
+        """
+        mqtt.publish_template(self.hass, "test/topic", "{{ 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")
+
+    def test_service_call_with_payload_doesnt_render_template(self):
+        """
+        If a 'payload' is provided then use that instead of 'payload_template'.
+        """
+        payload = "not a template"
+        payload_template = "a template"
+        # Call the service directly because the helper functions don't allow
+        # you to provide payload AND payload_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)
+        self.assertTrue(mqtt.MQTT_CLIENT.publish.called)
+        self.assertEqual(mqtt.MQTT_CLIENT.publish.call_args[0][1], payload)
+
+    def test_service_call_without_payload_or_payload_template(self):
+        """
+        If neither 'payload' or 'payload_template' is provided then fail.
+        """
+        # Call the service directly because the helper functions require you to
+        # provide a payload.
+        self.hass.services.call(mqtt.DOMAIN, mqtt.SERVICE_PUBLISH, {
+            mqtt.ATTR_TOPIC: "test/topic"
+        }, blocking=True)
+        self.assertFalse(mqtt.MQTT_CLIENT.publish.called)
+
     def test_subscribe_topic(self):
         mqtt.subscribe(self.hass, 'test-topic', self.record_calls)