From b69b9c60ececdab1d35621bdd7ffe5b8b9d59c44 Mon Sep 17 00:00:00 2001 From: jbouwh Date: Tue, 21 Dec 2021 16:52:59 +0000 Subject: [PATCH] Add encoding support for publishing --- homeassistant/components/mqtt/__init__.py | 37 ++++++++++++++++-- .../components/mqtt/alarm_control_panel.py | 13 ++++++- homeassistant/components/mqtt/climate.py | 6 ++- homeassistant/components/mqtt/cover.py | 17 +++++++-- homeassistant/components/mqtt/fan.py | 13 ++++++- homeassistant/components/mqtt/humidifier.py | 13 ++++++- .../components/mqtt/light/schema_basic.py | 12 +++++- homeassistant/components/mqtt/number.py | 13 ++++++- homeassistant/components/mqtt/select.py | 13 ++++++- tests/components/mqtt/test_init.py | 38 ++++++++++++++++++- 10 files changed, 154 insertions(+), 21 deletions(-) diff --git a/homeassistant/components/mqtt/__init__.py b/homeassistant/components/mqtt/__init__.py index 36f217e0bd9..2c29cf096e3 100644 --- a/homeassistant/components/mqtt/__init__.py +++ b/homeassistant/components/mqtt/__init__.py @@ -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, ) diff --git a/homeassistant/components/mqtt/alarm_control_panel.py b/homeassistant/components/mqtt/alarm_control_panel.py index 2a2c440e3ca..8609521b394 100644 --- a/homeassistant/components/mqtt/alarm_control_panel.py +++ b/homeassistant/components/mqtt/alarm_control_panel.py @@ -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): diff --git a/homeassistant/components/mqtt/climate.py b/homeassistant/components/mqtt/climate.py index 1edea7a62a6..78dbc584aca 100644 --- a/homeassistant/components/mqtt/climate.py +++ b/homeassistant/components/mqtt/climate.py @@ -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 diff --git a/homeassistant/components/mqtt/cover.py b/homeassistant/components/mqtt/cover.py index 9a950c8e44d..28b41adacf4 100644 --- a/homeassistant/components/mqtt/cover.py +++ b/homeassistant/components/mqtt/cover.py @@ -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) diff --git a/homeassistant/components/mqtt/fan.py b/homeassistant/components/mqtt/fan.py index 1837eb6a28f..e7566fc2d06 100644 --- a/homeassistant/components/mqtt/fan.py +++ b/homeassistant/components/mqtt/fan.py @@ -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(): diff --git a/homeassistant/components/mqtt/humidifier.py b/homeassistant/components/mqtt/humidifier.py index 74de601134c..4feb065ce9f 100644 --- a/homeassistant/components/mqtt/humidifier.py +++ b/homeassistant/components/mqtt/humidifier.py @@ -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(): diff --git a/homeassistant/components/mqtt/light/schema_basic.py b/homeassistant/components/mqtt/light/schema_basic.py index 7e6c34a33eb..e097518a921 100644 --- a/homeassistant/components/mqtt/light/schema_basic.py +++ b/homeassistant/components/mqtt/light/schema_basic.py @@ -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 diff --git a/homeassistant/components/mqtt/number.py b/homeassistant/components/mqtt/number.py index 11d714361a9..dc601aa825c 100644 --- a/homeassistant/components/mqtt/number.py +++ b/homeassistant/components/mqtt/number.py @@ -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), } diff --git a/homeassistant/components/mqtt/select.py b/homeassistant/components/mqtt/select.py index 9899860c86d..ab7819d1a0c 100644 --- a/homeassistant/components/mqtt/select.py +++ b/homeassistant/components/mqtt/select.py @@ -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), } diff --git a/tests/components/mqtt/test_init.py b/tests/components/mqtt/test_init.py index b2d468a66b8..819d0574c38 100644 --- a/tests/components/mqtt/test_init.py +++ b/tests/components/mqtt/test_init.py @@ -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."""