Add command_template to MQTT number platform (#58949)

This commit is contained in:
Jan Bouwhuis 2021-11-02 18:21:49 +01:00 committed by GitHub
parent 2df52a3bf6
commit 0c4863198e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 123 additions and 8 deletions

View File

@ -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],
)

View File

@ -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(