Add encoding support for publishing

pull/61937/head
jbouwh 2021-12-21 16:52:59 +00:00
parent cac2f8ac06
commit b69b9c60ec
10 changed files with 154 additions and 21 deletions

View File

@ -263,9 +263,11 @@ class MqttCommandTemplate:
*,
hass: HomeAssistant | None = None,
entity: Entity | None = None,
encoding: str | None = DEFAULT_ENCODING,
) -> None:
"""Instantiate a command template."""
self._attr_command_template = command_template
self._encoding = encoding
if command_template is None:
return
@ -299,8 +301,34 @@ class MqttCommandTemplate:
return payload
def _encode_outgoing_payload(
payload: PublishPayloadType, encoding: str | None = DEFAULT_ENCODING
) -> PublishPayloadType:
"""Ensure the correct encoding for the MQTT payload when publishing."""
if isinstance(payload, str):
if encoding is None:
_LOGGER.warning(
"Can't encode payload '%s' with no encoding set, encoding defaults to 'utf-8'",
payload,
)
return payload
if encoding == DEFAULT_ENCODING:
# No need to encode since UTF-8 is the default encoding for MQTT
return payload
try:
payload = payload.encode(encoding)
except (LookupError, UnicodeEncodeError):
_LOGGER.warning(
"Can't encode payload '%s' with encoding '%s', encoding defaults to 'utf-8'",
payload,
encoding,
)
return payload
if self._attr_command_template is None:
return value
return _encode_outgoing_payload(value, self._encoding)
values = {"value": value}
if self._entity:
@ -308,8 +336,11 @@ class MqttCommandTemplate:
values[ATTR_NAME] = self._entity.name
if variables is not None:
values.update(variables)
return _convert_outgoing_payload(
self._attr_command_template.async_render(values, parse_result=False)
return _encode_outgoing_payload(
_convert_outgoing_payload(
self._attr_command_template.async_render(values, parse_result=False),
),
self._encoding,
)

View File

@ -36,7 +36,14 @@ from homeassistant.helpers.typing import ConfigType
from . import PLATFORMS, MqttCommandTemplate, subscription
from .. import mqtt
from .const import CONF_COMMAND_TOPIC, CONF_QOS, CONF_RETAIN, CONF_STATE_TOPIC, DOMAIN
from .const import (
CONF_COMMAND_TOPIC,
CONF_ENCODING,
CONF_QOS,
CONF_RETAIN,
CONF_STATE_TOPIC,
DOMAIN,
)
from .debug_info import log_messages
from .mixins import MQTT_ENTITY_COMMON_SCHEMA, MqttEntity, async_setup_entry_helper
@ -151,7 +158,9 @@ class MqttAlarm(MqttEntity, alarm.AlarmControlPanelEntity):
if value_template is not None:
value_template.hass = self.hass
self._command_template = MqttCommandTemplate(
self._config[CONF_COMMAND_TEMPLATE], entity=self
self._config[CONF_COMMAND_TEMPLATE],
entity=self,
encoding=self._config.get(CONF_ENCODING) or None,
).async_render
async def _subscribe_topics(self):

View File

@ -54,7 +54,7 @@ from homeassistant.helpers.typing import ConfigType
from . import MQTT_BASE_PLATFORM_SCHEMA, PLATFORMS, MqttCommandTemplate, subscription
from .. import mqtt
from .const import CONF_QOS, CONF_RETAIN, DOMAIN
from .const import CONF_ENCODING, CONF_QOS, CONF_RETAIN, DOMAIN
from .debug_info import log_messages
from .mixins import MQTT_ENTITY_COMMON_SCHEMA, MqttEntity, async_setup_entry_helper
@ -378,7 +378,9 @@ class MqttClimate(MqttEntity, ClimateEntity):
command_templates = {}
for key in COMMAND_TEMPLATE_KEYS:
command_templates[key] = MqttCommandTemplate(
config.get(key), entity=self
config.get(key),
entity=self,
encoding=self._config.get(CONF_ENCODING) or None,
).async_render
self._command_templates = command_templates

View File

@ -38,7 +38,14 @@ from homeassistant.helpers.typing import ConfigType
from . import PLATFORMS, MqttCommandTemplate, subscription
from .. import mqtt
from .const import CONF_COMMAND_TOPIC, CONF_QOS, CONF_RETAIN, CONF_STATE_TOPIC, DOMAIN
from .const import (
CONF_COMMAND_TOPIC,
CONF_ENCODING,
CONF_QOS,
CONF_RETAIN,
CONF_STATE_TOPIC,
DOMAIN,
)
from .debug_info import log_messages
from .mixins import MQTT_ENTITY_COMMON_SCHEMA, MqttEntity, async_setup_entry_helper
@ -289,7 +296,9 @@ class MqttCover(MqttEntity, CoverEntity):
value_template.hass = self.hass
self._set_position_template = MqttCommandTemplate(
self._config.get(CONF_SET_POSITION_TEMPLATE), entity=self
self._config.get(CONF_SET_POSITION_TEMPLATE),
entity=self,
encoding=self._config.get(CONF_ENCODING),
).async_render
get_position_template = self._config.get(CONF_GET_POSITION_TEMPLATE)
@ -297,7 +306,9 @@ class MqttCover(MqttEntity, CoverEntity):
get_position_template.hass = self.hass
self._set_tilt_template = MqttCommandTemplate(
self._config.get(CONF_TILT_COMMAND_TEMPLATE), entity=self
self._config.get(CONF_TILT_COMMAND_TEMPLATE),
entity=self,
encoding=self._config.get(CONF_ENCODING) or None,
).async_render
tilt_status_template = self._config.get(CONF_TILT_STATUS_TEMPLATE)

View File

@ -38,7 +38,14 @@ from homeassistant.util.percentage import (
from . import PLATFORMS, MqttCommandTemplate, subscription
from .. import mqtt
from .const import CONF_COMMAND_TOPIC, CONF_QOS, CONF_RETAIN, CONF_STATE_TOPIC, DOMAIN
from .const import (
CONF_COMMAND_TOPIC,
CONF_ENCODING,
CONF_QOS,
CONF_RETAIN,
CONF_STATE_TOPIC,
DOMAIN,
)
from .debug_info import log_messages
from .mixins import MQTT_ENTITY_COMMON_SCHEMA, MqttEntity, async_setup_entry_helper
@ -334,7 +341,9 @@ class MqttFan(MqttEntity, FanEntity):
for key, tpl in self._command_templates.items():
self._command_templates[key] = MqttCommandTemplate(
tpl, entity=self
tpl,
entity=self,
encoding=self._config.get(CONF_ENCODING) or None,
).async_render
for key, tpl in self._value_templates.items():

View File

@ -29,7 +29,14 @@ from homeassistant.helpers.typing import ConfigType
from . import PLATFORMS, MqttCommandTemplate, subscription
from .. import mqtt
from .const import CONF_COMMAND_TOPIC, CONF_QOS, CONF_RETAIN, CONF_STATE_TOPIC, DOMAIN
from .const import (
CONF_COMMAND_TOPIC,
CONF_ENCODING,
CONF_QOS,
CONF_RETAIN,
CONF_STATE_TOPIC,
DOMAIN,
)
from .debug_info import log_messages
from .mixins import MQTT_ENTITY_COMMON_SCHEMA, MqttEntity, async_setup_entry_helper
@ -239,7 +246,9 @@ class MqttHumidifier(MqttEntity, HumidifierEntity):
for key, tpl in self._command_templates.items():
self._command_templates[key] = MqttCommandTemplate(
tpl, entity=self
tpl,
entity=self,
encoding=self._config.get(CONF_ENCODING) or None,
).async_render
for key, tpl in self._value_templates.items():

View File

@ -53,7 +53,13 @@ import homeassistant.util.color as color_util
from .. import MqttCommandTemplate, subscription
from ... import mqtt
from ..const import CONF_COMMAND_TOPIC, CONF_QOS, CONF_RETAIN, CONF_STATE_TOPIC
from ..const import (
CONF_COMMAND_TOPIC,
CONF_ENCODING,
CONF_QOS,
CONF_RETAIN,
CONF_STATE_TOPIC,
)
from ..debug_info import log_messages
from ..mixins import MQTT_ENTITY_COMMON_SCHEMA, MqttEntity
from .schema import MQTT_LIGHT_SCHEMA_SCHEMA
@ -331,7 +337,9 @@ class MqttLight(MqttEntity, LightEntity, RestoreEntity):
command_templates[key] = None
for key in COMMAND_TEMPLATE_KEYS & config.keys():
command_templates[key] = MqttCommandTemplate(
config[key], entity=self
config[key],
entity=self,
encoding=self._config.get(CONF_ENCODING) or None,
).async_render
self._command_templates = command_templates

View File

@ -27,7 +27,14 @@ from homeassistant.helpers.typing import ConfigType
from . import PLATFORMS, MqttCommandTemplate, subscription
from .. import mqtt
from .const import CONF_COMMAND_TOPIC, CONF_QOS, CONF_RETAIN, CONF_STATE_TOPIC, DOMAIN
from .const import (
CONF_COMMAND_TOPIC,
CONF_ENCODING,
CONF_QOS,
CONF_RETAIN,
CONF_STATE_TOPIC,
DOMAIN,
)
from .debug_info import log_messages
from .mixins import MQTT_ENTITY_COMMON_SCHEMA, MqttEntity, async_setup_entry_helper
@ -139,7 +146,9 @@ class MqttNumber(MqttEntity, NumberEntity, RestoreEntity):
self._templates = {
CONF_COMMAND_TEMPLATE: MqttCommandTemplate(
config.get(CONF_COMMAND_TEMPLATE), entity=self
config.get(CONF_COMMAND_TEMPLATE),
entity=self,
encoding=self._config.get(CONF_ENCODING) or None,
).async_render,
CONF_VALUE_TEMPLATE: config.get(CONF_VALUE_TEMPLATE),
}

View File

@ -15,7 +15,14 @@ from homeassistant.helpers.typing import ConfigType
from . import PLATFORMS, MqttCommandTemplate, subscription
from .. import mqtt
from .const import CONF_COMMAND_TOPIC, CONF_QOS, CONF_RETAIN, CONF_STATE_TOPIC, DOMAIN
from .const import (
CONF_COMMAND_TOPIC,
CONF_ENCODING,
CONF_QOS,
CONF_RETAIN,
CONF_STATE_TOPIC,
DOMAIN,
)
from .debug_info import log_messages
from .mixins import MQTT_ENTITY_COMMON_SCHEMA, MqttEntity, async_setup_entry_helper
@ -103,7 +110,9 @@ class MqttSelect(MqttEntity, SelectEntity, RestoreEntity):
self._templates = {
CONF_COMMAND_TEMPLATE: MqttCommandTemplate(
config.get(CONF_COMMAND_TEMPLATE), entity=self
config.get(CONF_COMMAND_TEMPLATE),
entity=self,
encoding=self._config.get(CONF_ENCODING) or None,
).async_render,
CONF_VALUE_TEMPLATE: config.get(CONF_VALUE_TEMPLATE),
}

View File

@ -158,7 +158,7 @@ async def test_publish(hass, mqtt_mock):
mqtt_mock.reset_mock()
async def test_convert_outgoing_payload(hass):
async def test_convert_outgoing_payload(hass, caplog):
"""Test the converting of outgoing MQTT payloads without template."""
command_template = mqtt.MqttCommandTemplate(None, hass=hass)
assert command_template.async_render(b"\xde\xad\xbe\xef") == b"\xde\xad\xbe\xef"
@ -174,6 +174,42 @@ async def test_convert_outgoing_payload(hass):
assert command_template.async_render(None) is None
# Test with different encoding
command_template = mqtt.MqttCommandTemplate(None, hass=hass, encoding="ascii")
assert (
command_template.async_render("I love Home Assistant")
== b"I love Home Assistant"
)
# Test with different encoding and non string payload
assert command_template.async_render(1234) == 1234
assert command_template.async_render(1234.56) == 1234.56
assert command_template.async_render(None) is None
# Tests with None encoding
command_template = mqtt.MqttCommandTemplate(None, hass=hass, encoding=None)
assert (
command_template.async_render("I love Home Assistant")
== "I love Home Assistant"
)
assert (
"Can't encode payload 'I love Home Assistant' with no encoding set, encoding defaults to 'utf-8'"
in caplog.text
)
# Test with invalid encoding
command_template = mqtt.MqttCommandTemplate(None, hass=hass, encoding="invalid")
assert (
command_template.async_render("I love Home Assistant")
== "I love Home Assistant"
)
assert "Can't encode payload '%s' with encoding 'invalid', encoding defaults to 'utf-8'"
async def test_command_template_value(hass):
"""Test the rendering of MQTT command template."""