Add command_template to MQTT number platform (#58949)

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)
+        self._templates = {
-        if value_template is not None:
+            CONF_COMMAND_TEMPLATE: config.get(CONF_COMMAND_TEMPLATE),
-            value_template.hass = self.hass
+            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
+            payload = self._templates[CONF_VALUE_TEMPLATE](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)
             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],
         )

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(