From 1bcd62cd32ff84d3a56d6bfd88589410e1eb9891 Mon Sep 17 00:00:00 2001 From: David Beitey Date: Wed, 17 Nov 2021 01:13:54 +1000 Subject: [PATCH] Add topic_template for mqtt.publish (#53743) Co-authored-by: Erik Montnemery Co-authored-by: Erik Montnemery --- homeassistant/components/mqtt/__init__.py | 50 +++++++++--- tests/components/mqtt/test_init.py | 97 +++++++++++++++++++++++ 2 files changed, 137 insertions(+), 10 deletions(-) diff --git a/homeassistant/components/mqtt/__init__.py b/homeassistant/components/mqtt/__init__.py index f42663ac4a8..efa14388a58 100644 --- a/homeassistant/components/mqtt/__init__.py +++ b/homeassistant/components/mqtt/__init__.py @@ -106,6 +106,7 @@ DEFAULT_KEEPALIVE = 60 DEFAULT_PROTOCOL = PROTOCOL_311 DEFAULT_TLS_PROTOCOL = "auto" +ATTR_TOPIC_TEMPLATE = "topic_template" ATTR_PAYLOAD_TEMPLATE = "payload_template" MAX_RECONNECT_WAIT = 300 # seconds @@ -220,15 +221,19 @@ MQTT_RW_PLATFORM_SCHEMA = MQTT_BASE_PLATFORM_SCHEMA.extend( ) # Service call validation schema -MQTT_PUBLISH_SCHEMA = vol.Schema( - { - vol.Required(ATTR_TOPIC): valid_publish_topic, - vol.Exclusive(ATTR_PAYLOAD, CONF_PAYLOAD): cv.string, - vol.Exclusive(ATTR_PAYLOAD_TEMPLATE, CONF_PAYLOAD): cv.string, - vol.Optional(ATTR_QOS, default=DEFAULT_QOS): _VALID_QOS_SCHEMA, - vol.Optional(ATTR_RETAIN, default=DEFAULT_RETAIN): cv.boolean, - }, - required=True, +MQTT_PUBLISH_SCHEMA = vol.All( + vol.Schema( + { + vol.Exclusive(ATTR_TOPIC, CONF_TOPIC): valid_publish_topic, + vol.Exclusive(ATTR_TOPIC_TEMPLATE, CONF_TOPIC): cv.string, + vol.Exclusive(ATTR_PAYLOAD, CONF_PAYLOAD): cv.string, + vol.Exclusive(ATTR_PAYLOAD_TEMPLATE, CONF_PAYLOAD): cv.string, + vol.Optional(ATTR_QOS, default=DEFAULT_QOS): _VALID_QOS_SCHEMA, + vol.Optional(ATTR_RETAIN, default=DEFAULT_RETAIN): cv.boolean, + }, + required=True, + ), + cv.has_at_least_one_key(ATTR_TOPIC, ATTR_TOPIC_TEMPLATE), ) @@ -450,11 +455,36 @@ async def async_setup_entry(hass, entry): async def async_publish_service(call: ServiceCall): """Handle MQTT publish service calls.""" - msg_topic: str = call.data[ATTR_TOPIC] + msg_topic = call.data.get(ATTR_TOPIC) + msg_topic_template = call.data.get(ATTR_TOPIC_TEMPLATE) payload = call.data.get(ATTR_PAYLOAD) payload_template = call.data.get(ATTR_PAYLOAD_TEMPLATE) qos: int = call.data[ATTR_QOS] retain: bool = call.data[ATTR_RETAIN] + if msg_topic_template is not None: + try: + rendered_topic = template.Template( + msg_topic_template, hass + ).async_render(parse_result=False) + msg_topic = valid_publish_topic(rendered_topic) + except (template.jinja2.TemplateError, TemplateError) as exc: + _LOGGER.error( + "Unable to publish: rendering topic template of %s " + "failed because %s", + msg_topic_template, + exc, + ) + return + except vol.Invalid as err: + _LOGGER.error( + "Unable to publish: topic template '%s' produced an " + "invalid topic '%s' after rendering (%s)", + msg_topic_template, + rendered_topic, + err, + ) + return + if payload_template is not None: try: payload = template.Template(payload_template, hass).async_render( diff --git a/tests/components/mqtt/test_init.py b/tests/components/mqtt/test_init.py index 26ceb583818..9b862e38a7c 100644 --- a/tests/components/mqtt/test_init.py +++ b/tests/components/mqtt/test_init.py @@ -150,6 +150,103 @@ async def test_service_call_without_topic_does_not_publish(hass, mqtt_mock): assert not mqtt_mock.async_publish.called +async def test_service_call_with_topic_and_topic_template_does_not_publish( + hass, mqtt_mock +): + """Test the service call with topic/topic template. + + If both 'topic' and 'topic_template' are provided then fail. + """ + topic = "test/topic" + topic_template = "test/{{ 'topic' }}" + with pytest.raises(vol.Invalid): + await hass.services.async_call( + mqtt.DOMAIN, + mqtt.SERVICE_PUBLISH, + { + mqtt.ATTR_TOPIC: topic, + mqtt.ATTR_TOPIC_TEMPLATE: topic_template, + mqtt.ATTR_PAYLOAD: "payload", + }, + blocking=True, + ) + assert not mqtt_mock.async_publish.called + + +async def test_service_call_with_invalid_topic_template_does_not_publish( + hass, mqtt_mock +): + """Test the service call with a problematic topic template.""" + await hass.services.async_call( + mqtt.DOMAIN, + mqtt.SERVICE_PUBLISH, + { + mqtt.ATTR_TOPIC_TEMPLATE: "test/{{ 1 | no_such_filter }}", + mqtt.ATTR_PAYLOAD: "payload", + }, + blocking=True, + ) + assert not mqtt_mock.async_publish.called + + +async def test_service_call_with_template_topic_renders_template(hass, mqtt_mock): + """Test the service call with rendered topic template. + + If 'topic_template' is provided and 'topic' is not, then render it. + """ + await hass.services.async_call( + mqtt.DOMAIN, + mqtt.SERVICE_PUBLISH, + { + mqtt.ATTR_TOPIC_TEMPLATE: "test/{{ 1+1 }}", + mqtt.ATTR_PAYLOAD: "payload", + }, + blocking=True, + ) + assert mqtt_mock.async_publish.called + assert mqtt_mock.async_publish.call_args[0][0] == "test/2" + + +async def test_service_call_with_template_topic_renders_invalid_topic(hass, mqtt_mock): + """Test the service call with rendered, invalid topic template. + + If a wildcard topic is rendered, then fail. + """ + await hass.services.async_call( + mqtt.DOMAIN, + mqtt.SERVICE_PUBLISH, + { + mqtt.ATTR_TOPIC_TEMPLATE: "test/{{ '+' if True else 'topic' }}/topic", + mqtt.ATTR_PAYLOAD: "payload", + }, + blocking=True, + ) + assert not mqtt_mock.async_publish.called + + +async def test_service_call_with_invalid_rendered_template_topic_doesnt_render_template( + hass, mqtt_mock +): + """Test the service call with unrendered template. + + If both 'payload' and 'payload_template' are provided then fail. + """ + payload = "not a template" + payload_template = "a template" + with pytest.raises(vol.Invalid): + await hass.services.async_call( + mqtt.DOMAIN, + mqtt.SERVICE_PUBLISH, + { + mqtt.ATTR_TOPIC: "test/topic", + mqtt.ATTR_PAYLOAD: payload, + mqtt.ATTR_PAYLOAD_TEMPLATE: payload_template, + }, + blocking=True, + ) + assert not mqtt_mock.async_publish.called + + async def test_service_call_with_template_payload_renders_template(hass, mqtt_mock): """Test the service call with rendered template.