Add MQTT availability template and encoding (#60470)

* Add MQTT availability template and encoding

* use generic encoding field

* pylint and cleanup

* remove additional topic check

homeassistant/components/mqtt/abbreviations.py

@@ -10,6 +10,7 @@ ABBREVIATIONS = {
     "avty": "availability",
     "avty_mode": "availability_mode",
     "avty_t": "availability_topic",
+    "avty_tpl": "availability_template",
     "away_mode_cmd_t": "away_mode_command_topic",
     "away_mode_stat_tpl": "away_mode_state_template",
     "away_mode_stat_t": "away_mode_state_topic",

homeassistant/components/mqtt/mixins.py

@@ -21,6 +21,7 @@ from homeassistant.const import (
     CONF_ICON,
     CONF_NAME,
     CONF_UNIQUE_ID,
+    CONF_VALUE_TEMPLATE,
 )
 from homeassistant.core import callback
 from homeassistant.helpers import config_validation as cv
@@ -42,6 +43,7 @@ from .const import (
     ATTR_DISCOVERY_PAYLOAD,
     ATTR_DISCOVERY_TOPIC,
     CONF_AVAILABILITY,
+    CONF_ENCODING,
     CONF_QOS,
     CONF_TOPIC,
     DEFAULT_PAYLOAD_AVAILABLE,
@@ -71,6 +73,7 @@ AVAILABILITY_LATEST = "latest"
 AVAILABILITY_MODES = [AVAILABILITY_ALL, AVAILABILITY_ANY, AVAILABILITY_LATEST]
 
 CONF_AVAILABILITY_MODE = "availability_mode"
+CONF_AVAILABILITY_TEMPLATE = "availability_template"
 CONF_AVAILABILITY_TOPIC = "availability_topic"
 CONF_ENABLED_BY_DEFAULT = "enabled_by_default"
 CONF_PAYLOAD_AVAILABLE = "payload_available"
@@ -112,6 +115,7 @@ MQTT_ATTRIBUTES_BLOCKED = {
 MQTT_AVAILABILITY_SINGLE_SCHEMA = vol.Schema(
     {
         vol.Exclusive(CONF_AVAILABILITY_TOPIC, "availability"): valid_subscribe_topic,
+        vol.Optional(CONF_AVAILABILITY_TEMPLATE): cv.template,
         vol.Optional(
             CONF_PAYLOAD_AVAILABLE, default=DEFAULT_PAYLOAD_AVAILABLE
         ): cv.string,
@@ -138,6 +142,7 @@ MQTT_AVAILABILITY_LIST_SCHEMA = vol.Schema(
                         CONF_PAYLOAD_NOT_AVAILABLE,
                         default=DEFAULT_PAYLOAD_NOT_AVAILABLE,
                     ): cv.string,
+                    vol.Optional(CONF_VALUE_TEMPLATE): cv.template,
                 }
             ],
         ),
@@ -335,6 +340,7 @@ class MqttAvailability(Entity):
             self._avail_topics[config[CONF_AVAILABILITY_TOPIC]] = {
                 CONF_PAYLOAD_AVAILABLE: config[CONF_PAYLOAD_AVAILABLE],
                 CONF_PAYLOAD_NOT_AVAILABLE: config[CONF_PAYLOAD_NOT_AVAILABLE],
+                CONF_AVAILABILITY_TEMPLATE: config.get(CONF_AVAILABILITY_TEMPLATE),
             }
 
         if CONF_AVAILABILITY in config:
@@ -342,8 +348,22 @@ class MqttAvailability(Entity):
                 self._avail_topics[avail[CONF_TOPIC]] = {
                     CONF_PAYLOAD_AVAILABLE: avail[CONF_PAYLOAD_AVAILABLE],
                     CONF_PAYLOAD_NOT_AVAILABLE: avail[CONF_PAYLOAD_NOT_AVAILABLE],
+                    CONF_AVAILABILITY_TEMPLATE: avail.get(CONF_VALUE_TEMPLATE),
                 }
 
+        for (
+            topic,  # pylint: disable=unused-variable
+            avail_topic_conf,
+        ) in self._avail_topics.items():
+            tpl = avail_topic_conf[CONF_AVAILABILITY_TEMPLATE]
+            if tpl is None:
+                avail_topic_conf[CONF_AVAILABILITY_TEMPLATE] = lambda value: value
+            else:
+                tpl.hass = self.hass
+                avail_topic_conf[
+                    CONF_AVAILABILITY_TEMPLATE
+                ] = tpl.async_render_with_possible_json_value
+
         self._avail_config = config
 
     async def _availability_subscribe_topics(self):
@@ -354,10 +374,11 @@ class MqttAvailability(Entity):
         def availability_message_received(msg: ReceiveMessage) -> None:
             """Handle a new received MQTT availability message."""
             topic = msg.topic
-            if msg.payload == self._avail_topics[topic][CONF_PAYLOAD_AVAILABLE]:
+            payload = self._avail_topics[topic][CONF_AVAILABILITY_TEMPLATE](msg.payload)
+            if payload == self._avail_topics[topic][CONF_PAYLOAD_AVAILABLE]:
                 self._available[topic] = True
                 self._available_latest = True
-            elif msg.payload == self._avail_topics[topic][CONF_PAYLOAD_NOT_AVAILABLE]:
+            elif payload == self._avail_topics[topic][CONF_PAYLOAD_NOT_AVAILABLE]:
                 self._available[topic] = False
                 self._available_latest = False
 
@@ -372,6 +393,7 @@ class MqttAvailability(Entity):
                 "topic": topic,
                 "msg_callback": availability_message_received,
                 "qos": self._avail_config[CONF_QOS],
+                "encoding": self._avail_config[CONF_ENCODING] or None,
             }
             for topic in self._avail_topics
         }

tests/components/mqtt/test_discovery.py

@@ -737,6 +737,100 @@ async def test_discovery_expansion_3(hass, mqtt_mock, caplog):
     )
 
 
+async def test_discovery_expansion_without_encoding_and_value_template_1(
+    hass, mqtt_mock, caplog
+):
+    """Test expansion of raw availability payload with a template as list."""
+    data = (
+        '{ "~": "some/base/topic",'
+        '  "name": "DiscoveryExpansionTest1",'
+        '  "stat_t": "test_topic/~",'
+        '  "cmd_t": "~/test_topic",'
+        '  "encoding":"",'
+        '  "availability": [{'
+        '    "topic":"~/avail_item1",'
+        '    "payload_available": "1",'
+        '    "payload_not_available": "0",'
+        '    "value_template":"{{ value | bitwise_and(1) }}"'
+        "  }],"
+        '  "dev":{'
+        '    "ids":["5706DF"],'
+        '    "name":"DiscoveryExpansionTest1 Device",'
+        '    "mdl":"Generic",'
+        '    "sw":"1.2.3.4",'
+        '    "mf":"None",'
+        '    "sa":"default_area"'
+        "  }"
+        "}"
+    )
+
+    async_fire_mqtt_message(hass, "homeassistant/switch/bla/config", data)
+    await hass.async_block_till_done()
+
+    state = hass.states.get("switch.DiscoveryExpansionTest1")
+    assert state.state == STATE_UNAVAILABLE
+
+    async_fire_mqtt_message(hass, "some/base/topic/avail_item1", b"\x01")
+    await hass.async_block_till_done()
+
+    state = hass.states.get("switch.DiscoveryExpansionTest1")
+    assert state is not None
+    assert state.name == "DiscoveryExpansionTest1"
+    assert ("switch", "bla") in hass.data[ALREADY_DISCOVERED]
+    assert state.state == STATE_OFF
+
+    async_fire_mqtt_message(hass, "some/base/topic/avail_item1", b"\x00")
+
+    state = hass.states.get("switch.DiscoveryExpansionTest1")
+    assert state.state == STATE_UNAVAILABLE
+
+
+async def test_discovery_expansion_without_encoding_and_value_template_2(
+    hass, mqtt_mock, caplog
+):
+    """Test expansion of raw availability payload with a template directly."""
+    data = (
+        '{ "~": "some/base/topic",'
+        '  "name": "DiscoveryExpansionTest1",'
+        '  "stat_t": "test_topic/~",'
+        '  "cmd_t": "~/test_topic",'
+        '  "availability_topic":"~/avail_item1",'
+        '  "payload_available": "1",'
+        '  "payload_not_available": "0",'
+        '  "encoding":"",'
+        '  "availability_template":"{{ value | bitwise_and(1) }}",'
+        '  "dev":{'
+        '    "ids":["5706DF"],'
+        '    "name":"DiscoveryExpansionTest1 Device",'
+        '    "mdl":"Generic",'
+        '    "sw":"1.2.3.4",'
+        '    "mf":"None",'
+        '    "sa":"default_area"'
+        "  }"
+        "}"
+    )
+
+    async_fire_mqtt_message(hass, "homeassistant/switch/bla/config", data)
+    await hass.async_block_till_done()
+
+    state = hass.states.get("switch.DiscoveryExpansionTest1")
+    assert state.state == STATE_UNAVAILABLE
+
+    async_fire_mqtt_message(hass, "some/base/topic/avail_item1", b"\x01")
+    await hass.async_block_till_done()
+
+    state = hass.states.get("switch.DiscoveryExpansionTest1")
+    assert state is not None
+    assert state.name == "DiscoveryExpansionTest1"
+    assert ("switch", "bla") in hass.data[ALREADY_DISCOVERED]
+    assert state.state == STATE_OFF
+
+    async_fire_mqtt_message(hass, "some/base/topic/avail_item1", b"\x00")
+
+    state = hass.states.get("switch.DiscoveryExpansionTest1")
+    assert state.state == STATE_UNAVAILABLE
+
+
 ABBREVIATIONS_WHITE_LIST = [
     # MQTT client/server/trigger settings
     "CONF_BIRTH_MESSAGE",