From 0c4863198e8c6f67f9e358f7cffc95e9fcfc63b2 Mon Sep 17 00:00:00 2001 From: Jan Bouwhuis Date: Tue, 2 Nov 2021 18:21:49 +0100 Subject: [PATCH] Add command_template to MQTT number platform (#58949) --- homeassistant/components/mqtt/number.py | 24 ++++-- tests/components/mqtt/test_number.py | 107 ++++++++++++++++++++++++ 2 files changed, 123 insertions(+), 8 deletions(-) diff --git a/homeassistant/components/mqtt/number.py b/homeassistant/components/mqtt/number.py index f39ca870505..6d0163abe32 100644 --- a/homeassistant/components/mqtt/number.py +++ b/homeassistant/components/mqtt/number.py @@ -31,6 +31,8 @@ from .const import CONF_COMMAND_TOPIC, CONF_QOS, CONF_RETAIN, CONF_STATE_TOPIC, from .debug_info import log_messages from .mixins import MQTT_ENTITY_COMMON_SCHEMA, MqttEntity, async_setup_entry_helper +CONF_COMMAND_TEMPLATE = "command_template" + _LOGGER = logging.getLogger(__name__) CONF_MIN = "min" @@ -61,6 +63,7 @@ def validate_config(config): _PLATFORM_SCHEMA_BASE = mqtt.MQTT_RW_PLATFORM_SCHEMA.extend( { + vol.Optional(CONF_COMMAND_TEMPLATE): cv.template, vol.Optional(CONF_MAX, default=DEFAULT_MAX_VALUE): vol.Coerce(float), vol.Optional(CONF_MIN, default=DEFAULT_MIN_VALUE): vol.Coerce(float), vol.Optional(CONF_NAME, default=DEFAULT_NAME): cv.string, @@ -133,9 +136,16 @@ class MqttNumber(MqttEntity, NumberEntity, RestoreEntity): """(Re)Setup the entity.""" self._optimistic = config[CONF_OPTIMISTIC] - value_template = self._config.get(CONF_VALUE_TEMPLATE) - if value_template is not None: - value_template.hass = self.hass + self._templates = { + CONF_COMMAND_TEMPLATE: config.get(CONF_COMMAND_TEMPLATE), + CONF_VALUE_TEMPLATE: config.get(CONF_VALUE_TEMPLATE), + } + for key, tpl in self._templates.items(): + if tpl is None: + self._templates[key] = lambda value: value + else: + tpl.hass = self.hass + self._templates[key] = tpl.async_render_with_possible_json_value async def _subscribe_topics(self): """(Re)Subscribe to topics.""" @@ -144,10 +154,7 @@ class MqttNumber(MqttEntity, NumberEntity, RestoreEntity): @log_messages(self.hass, self.entity_id) def message_received(msg): """Handle new MQTT messages.""" - payload = msg.payload - value_template = self._config.get(CONF_VALUE_TEMPLATE) - if value_template is not None: - payload = value_template.async_render_with_possible_json_value(payload) + payload = self._templates[CONF_VALUE_TEMPLATE](msg.payload) try: if payload == self._config[CONF_PAYLOAD_RESET]: num_value = None @@ -224,6 +231,7 @@ class MqttNumber(MqttEntity, NumberEntity, RestoreEntity): if value.is_integer(): current_number = int(value) + payload = self._templates[CONF_COMMAND_TEMPLATE](current_number) if self._optimistic: self._current_number = current_number @@ -232,7 +240,7 @@ class MqttNumber(MqttEntity, NumberEntity, RestoreEntity): await mqtt.async_publish( self.hass, self._config[CONF_COMMAND_TOPIC], - current_number, + payload, self._config[CONF_QOS], self._config[CONF_RETAIN], ) diff --git a/tests/components/mqtt/test_number.py b/tests/components/mqtt/test_number.py index 3b8b370eff8..797a7b894fc 100644 --- a/tests/components/mqtt/test_number.py +++ b/tests/components/mqtt/test_number.py @@ -209,6 +209,76 @@ async def test_run_number_service_optimistic(hass, mqtt_mock): assert state.state == "42.1" +async def test_run_number_service_optimistic_with_command_template(hass, mqtt_mock): + """Test that set_value service works in optimistic mode and with a command_template.""" + topic = "test/number" + + fake_state = ha.State("switch.test", "3") + + with patch( + "homeassistant.helpers.restore_state.RestoreEntity.async_get_last_state", + return_value=fake_state, + ): + assert await async_setup_component( + hass, + number.DOMAIN, + { + "number": { + "platform": "mqtt", + "command_topic": topic, + "name": "Test Number", + "command_template": '{"number": {{ value }} }', + } + }, + ) + await hass.async_block_till_done() + + state = hass.states.get("number.test_number") + assert state.state == "3" + assert state.attributes.get(ATTR_ASSUMED_STATE) + + # Integer + await hass.services.async_call( + NUMBER_DOMAIN, + SERVICE_SET_VALUE, + {ATTR_ENTITY_ID: "number.test_number", ATTR_VALUE: 30}, + blocking=True, + ) + + mqtt_mock.async_publish.assert_called_once_with(topic, '{"number": 30 }', 0, False) + mqtt_mock.async_publish.reset_mock() + state = hass.states.get("number.test_number") + assert state.state == "30" + + # Float with no decimal -> integer + await hass.services.async_call( + NUMBER_DOMAIN, + SERVICE_SET_VALUE, + {ATTR_ENTITY_ID: "number.test_number", ATTR_VALUE: 42.0}, + blocking=True, + ) + + mqtt_mock.async_publish.assert_called_once_with(topic, '{"number": 42 }', 0, False) + mqtt_mock.async_publish.reset_mock() + state = hass.states.get("number.test_number") + assert state.state == "42" + + # Float with decimal -> float + await hass.services.async_call( + NUMBER_DOMAIN, + SERVICE_SET_VALUE, + {ATTR_ENTITY_ID: "number.test_number", ATTR_VALUE: 42.1}, + blocking=True, + ) + + mqtt_mock.async_publish.assert_called_once_with( + topic, '{"number": 42.1 }', 0, False + ) + mqtt_mock.async_publish.reset_mock() + state = hass.states.get("number.test_number") + assert state.state == "42.1" + + async def test_run_number_service(hass, mqtt_mock): """Test that set_value service works in non optimistic mode.""" cmd_topic = "test/number/set" @@ -243,6 +313,43 @@ async def test_run_number_service(hass, mqtt_mock): assert state.state == "32" +async def test_run_number_service_with_command_template(hass, mqtt_mock): + """Test that set_value service works in non optimistic mode and with a command_template.""" + cmd_topic = "test/number/set" + state_topic = "test/number" + + assert await async_setup_component( + hass, + number.DOMAIN, + { + "number": { + "platform": "mqtt", + "command_topic": cmd_topic, + "state_topic": state_topic, + "name": "Test Number", + "command_template": '{"number": {{ value }} }', + } + }, + ) + await hass.async_block_till_done() + + async_fire_mqtt_message(hass, state_topic, "32") + state = hass.states.get("number.test_number") + assert state.state == "32" + + await hass.services.async_call( + NUMBER_DOMAIN, + SERVICE_SET_VALUE, + {ATTR_ENTITY_ID: "number.test_number", ATTR_VALUE: 30}, + blocking=True, + ) + mqtt_mock.async_publish.assert_called_once_with( + cmd_topic, '{"number": 30 }', 0, False + ) + state = hass.states.get("number.test_number") + assert state.state == "32" + + async def test_availability_when_connection_lost(hass, mqtt_mock): """Test availability after MQTT disconnection.""" await help_test_availability_when_connection_lost(