Add topic_template for mqtt.publish (#53743)

Co-authored-by: Erik Montnemery <erik@montnemery.com>

Co-authored-by: Erik Montnemery <erik@montnemery.com>
This commit is contained in:
David Beitey 2021-11-17 01:13:54 +10:00 committed by GitHub
parent 8a02d87a17
commit 1bcd62cd32
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 137 additions and 10 deletions

View file

@ -106,6 +106,7 @@ DEFAULT_KEEPALIVE = 60
DEFAULT_PROTOCOL = PROTOCOL_311 DEFAULT_PROTOCOL = PROTOCOL_311
DEFAULT_TLS_PROTOCOL = "auto" DEFAULT_TLS_PROTOCOL = "auto"
ATTR_TOPIC_TEMPLATE = "topic_template"
ATTR_PAYLOAD_TEMPLATE = "payload_template" ATTR_PAYLOAD_TEMPLATE = "payload_template"
MAX_RECONNECT_WAIT = 300 # seconds MAX_RECONNECT_WAIT = 300 # seconds
@ -220,15 +221,19 @@ MQTT_RW_PLATFORM_SCHEMA = MQTT_BASE_PLATFORM_SCHEMA.extend(
) )
# Service call validation schema # Service call validation schema
MQTT_PUBLISH_SCHEMA = vol.Schema( MQTT_PUBLISH_SCHEMA = vol.All(
{ vol.Schema(
vol.Required(ATTR_TOPIC): valid_publish_topic, {
vol.Exclusive(ATTR_PAYLOAD, CONF_PAYLOAD): cv.string, vol.Exclusive(ATTR_TOPIC, CONF_TOPIC): valid_publish_topic,
vol.Exclusive(ATTR_PAYLOAD_TEMPLATE, CONF_PAYLOAD): cv.string, vol.Exclusive(ATTR_TOPIC_TEMPLATE, CONF_TOPIC): cv.string,
vol.Optional(ATTR_QOS, default=DEFAULT_QOS): _VALID_QOS_SCHEMA, vol.Exclusive(ATTR_PAYLOAD, CONF_PAYLOAD): cv.string,
vol.Optional(ATTR_RETAIN, default=DEFAULT_RETAIN): cv.boolean, vol.Exclusive(ATTR_PAYLOAD_TEMPLATE, CONF_PAYLOAD): cv.string,
}, vol.Optional(ATTR_QOS, default=DEFAULT_QOS): _VALID_QOS_SCHEMA,
required=True, 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): async def async_publish_service(call: ServiceCall):
"""Handle MQTT publish service calls.""" """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 = call.data.get(ATTR_PAYLOAD)
payload_template = call.data.get(ATTR_PAYLOAD_TEMPLATE) payload_template = call.data.get(ATTR_PAYLOAD_TEMPLATE)
qos: int = call.data[ATTR_QOS] qos: int = call.data[ATTR_QOS]
retain: bool = call.data[ATTR_RETAIN] 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: if payload_template is not None:
try: try:
payload = template.Template(payload_template, hass).async_render( payload = template.Template(payload_template, hass).async_render(

View file

@ -150,6 +150,103 @@ async def test_service_call_without_topic_does_not_publish(hass, mqtt_mock):
assert not mqtt_mock.async_publish.called 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): async def test_service_call_with_template_payload_renders_template(hass, mqtt_mock):
"""Test the service call with rendered template. """Test the service call with rendered template.