Add availability support for MQTT subentries (#138673)

* Add availability support for MQTT subentries

* Update homeassistant/components/mqtt/config_flow.py

Co-authored-by: Erik Montnemery <erik@montnemery.com>

* Update homeassistant/components/mqtt/config_flow.py

Co-authored-by: Erik Montnemery <erik@montnemery.com>

* Update homeassistant/components/mqtt/config_flow.py

Co-authored-by: Erik Montnemery <erik@montnemery.com>

* Update homeassistant/components/mqtt/strings.json

Co-authored-by: Erik Montnemery <erik@montnemery.com>

---------

Co-authored-by: Erik Montnemery <erik@montnemery.com>

homeassistant/components/mqtt/config_flow.py

@@ -88,6 +88,8 @@ from .const import (
     ATTR_QOS,
     ATTR_RETAIN,
     ATTR_TOPIC,
+    CONF_AVAILABILITY_TEMPLATE,
+    CONF_AVAILABILITY_TOPIC,
     CONF_BIRTH_MESSAGE,
     CONF_BROKER,
     CONF_CERTIFICATE,
@@ -98,6 +100,8 @@ from .const import (
     CONF_DISCOVERY_PREFIX,
     CONF_ENTITY_PICTURE,
     CONF_KEEPALIVE,
+    CONF_PAYLOAD_AVAILABLE,
+    CONF_PAYLOAD_NOT_AVAILABLE,
     CONF_QOS,
     CONF_RETAIN,
     CONF_TLS_INSECURE,
@@ -111,6 +115,8 @@ from .const import (
     DEFAULT_DISCOVERY,
     DEFAULT_ENCODING,
     DEFAULT_KEEPALIVE,
+    DEFAULT_PAYLOAD_AVAILABLE,
+    DEFAULT_PAYLOAD_NOT_AVAILABLE,
     DEFAULT_PORT,
     DEFAULT_PREFIX,
     DEFAULT_PROTOCOL,
@@ -123,13 +129,15 @@ from .const import (
     TRANSPORT_WEBSOCKETS,
     Platform,
 )
-from .models import MqttDeviceData, MqttSubentryData
+from .models import MqttAvailabilityData, MqttDeviceData, MqttSubentryData
 from .util import (
     async_create_certificate_temp_files,
     get_file_path,
     valid_birth_will,
     valid_publish_topic,
     valid_qos_schema,
+    valid_subscribe_topic,
+    valid_subscribe_topic_template,
 )
 
 _LOGGER = logging.getLogger(__name__)
@@ -220,6 +228,19 @@ SUBENTRY_PLATFORM_SELECTOR = SelectSelector(
 
 TEMPLATE_SELECTOR = TemplateSelector(TemplateSelectorConfig())
 
+SUBENTRY_AVAILABILITY_SCHEMA = vol.Schema(
+    {
+        vol.Optional(CONF_AVAILABILITY_TOPIC): TEXT_SELECTOR,
+        vol.Optional(CONF_AVAILABILITY_TEMPLATE): TEMPLATE_SELECTOR,
+        vol.Optional(
+            CONF_PAYLOAD_AVAILABLE, default=DEFAULT_PAYLOAD_AVAILABLE
+        ): TEXT_SELECTOR,
+        vol.Optional(
+            CONF_PAYLOAD_NOT_AVAILABLE, default=DEFAULT_PAYLOAD_NOT_AVAILABLE
+        ): TEXT_SELECTOR,
+    }
+)
+
 
 @dataclass(frozen=True)
 class PlatformField:
@@ -1085,6 +1106,44 @@ class MQTTSubentryFlowHandler(ConfigSubentryFlow):
             },
         )
 
+    async def async_step_availability(
+        self, user_input: dict[str, Any] | None = None
+    ) -> SubentryFlowResult:
+        """Configure availability options."""
+        errors: dict[str, str] = {}
+        validate_field(
+            "availability_topic",
+            valid_subscribe_topic,
+            user_input,
+            errors,
+            "invalid_subscribe_topic",
+        )
+        validate_field(
+            "availability_template",
+            valid_subscribe_topic_template,
+            user_input,
+            errors,
+            "invalid_template",
+        )
+        if not errors and user_input is not None:
+            self._subentry_data.setdefault("availability", MqttAvailabilityData())
+            self._subentry_data["availability"] = cast(MqttAvailabilityData, user_input)
+            return await self.async_step_summary_menu()
+
+        data_schema = SUBENTRY_AVAILABILITY_SCHEMA
+        data_schema = self.add_suggested_values_to_schema(
+            data_schema,
+            dict(self._subentry_data.setdefault("availability", {}))
+            if self.source == SOURCE_RECONFIGURE
+            else user_input,
+        )
+        return self.async_show_form(
+            step_id="availability",
+            data_schema=data_schema,
+            errors=errors,
+            last_step=False,
+        )
+
     async def async_step_summary_menu(
         self, user_input: dict[str, Any] | None = None
     ) -> SubentryFlowResult:
@@ -1101,7 +1160,7 @@ class MQTTSubentryFlowHandler(ConfigSubentryFlow):
         ]
         if len(self._subentry_data["components"]) > 1:
             menu_options.append("delete_entity")
-        menu_options.append("device")
+        menu_options.extend(["device", "availability"])
         if self._subentry_data != self._get_reconfigure_subentry().data:
             menu_options.append("save_changes")
         return self.async_show_menu(

homeassistant/components/mqtt/entity.py

@@ -297,6 +297,7 @@ def async_setup_entity_entry_helper(
         # process subentry entity setup
         for config_subentry_id, subentry in entry.subentries.items():
             subentry_data = cast(MqttSubentryData, subentry.data)
+            availability_config = subentry_data.get("availability", {})
             subentry_entities: list[Entity] = []
             device_config = subentry_data["device"].copy()
             device_config["identifiers"] = config_subentry_id
@@ -309,6 +310,7 @@ def async_setup_entity_entry_helper(
                 )
                 component_config[CONF_DEVICE] = device_config
                 component_config.pop("platform")
+                component_config.update(availability_config)
 
                 try:
                     config = platform_schema_modern(component_config)

homeassistant/components/mqtt/models.py

@@ -432,11 +432,21 @@ class MqttDeviceData(TypedDict, total=False):
     model_id: str
 
 
-class MqttSubentryData(TypedDict):
+class MqttAvailabilityData(TypedDict, total=False):
+    """Hold the availability configuration for a device."""
+
+    availability_topic: str
+    availability_template: str
+    payload_available: str
+    payload_not_available: str
+
+
+class MqttSubentryData(TypedDict, total=False):
     """Hold the data for a MQTT subentry."""
 
     device: MqttDeviceData
     components: dict[str, dict[str, Any]]
+    availability: MqttAvailabilityData
 
 
 DATA_MQTT: HassKey[MqttData] = HassKey("mqtt")

homeassistant/components/mqtt/strings.json

@@ -116,6 +116,22 @@
       },
       "entry_type": "MQTT Device",
       "step": {
+        "availability": {
+          "title": "Availability options",
+          "description": "The availability feature allows a device to report it's availability.",
+          "data": {
+            "availability_topic": "Availability topic",
+            "availability_template": "Availability template",
+            "payload_available": "Payload available",
+            "payload_not_available": "Payload not available"
+          },
+          "data_description": {
+            "availability_topic": "Topic to receive the availabillity payload on",
+            "availability_template": "A [template](https://www.home-assistant.io/docs/configuration/templating/#using-templates-with-the-mqtt-integration) to render the availability payload received on the availability topic",
+            "payload_available": "The payload that indicates the device is available (defaults to 'online')",
+            "payload_not_available": "The payload that indicates the device is not available (defaults to 'offline')"
+          }
+        },
         "device": {
           "title": "Configure MQTT device details",
           "description": "Enter the MQTT device details:",
@@ -143,6 +159,7 @@
             "entity": "Add another entity to \"{mqtt_device}\"",
             "update_entity": "Update entity properties",
             "delete_entity": "Delete an entity",
+            "availability": "Configure availability",
             "device": "Update device properties",
             "save_changes": "Save changes"
           }

tests/components/mqtt/common.py

@@ -119,6 +119,15 @@ MOCK_SUBENTRY_NOTIFY_BAD_SCHEMA = {
     },
 }
 
+MOCK_SUBENTRY_AVAILABILITY_DATA = {
+    "availability": {
+        "availability_topic": "test/availability",
+        "availability_template": "{{ value_json.availability }}",
+        "payload_available": "online",
+        "payload_not_available": "offline",
+    }
+}
+
 MOCK_NOTIFY_SUBENTRY_DATA_MULTI = {
     "device": {
         "name": "Milk notifier",
@@ -129,7 +138,7 @@ MOCK_NOTIFY_SUBENTRY_DATA_MULTI = {
         "configuration_url": "https://example.com",
     },
     "components": MOCK_SUBENTRY_NOTIFY_COMPONENT1 | MOCK_SUBENTRY_NOTIFY_COMPONENT2,
-}
+} | MOCK_SUBENTRY_AVAILABILITY_DATA
 
 MOCK_NOTIFY_SUBENTRY_DATA_SINGLE = {
     "device": {
@@ -177,7 +186,7 @@ MOCK_SUBENTRY_DATA_SET_MIX = {
     "components": MOCK_SUBENTRY_NOTIFY_COMPONENT1
     | MOCK_SUBENTRY_NOTIFY_COMPONENT2
     | MOCK_SUBENTRY_LIGHT_COMPONENT,
-}
+} | MOCK_SUBENTRY_AVAILABILITY_DATA
 _SENTINEL = object()
 
 DISCOVERY_COUNT = sum(len(discovery_topic) for discovery_topic in MQTT.values())

tests/components/mqtt/test_config_flow.py

@@ -2821,6 +2821,7 @@ async def test_subentry_reconfigure_remove_entity(
         "update_entity",
         "delete_entity",
         "device",
+        "availability",
     ]
 
     # assert we can delete an entity
@@ -2849,6 +2850,7 @@ async def test_subentry_reconfigure_remove_entity(
         "entity",
         "update_entity",
         "device",
+        "availability",
         "save_changes",
     ]
 
@@ -2938,6 +2940,7 @@ async def test_subentry_reconfigure_edit_entity_multi_entitites(
         "update_entity",
         "delete_entity",
         "device",
+        "availability",
     ]
 
     # assert we can update an entity
@@ -3061,6 +3064,7 @@ async def test_subentry_reconfigure_edit_entity_single_entity(
         "entity",
         "update_entity",
         "device",
+        "availability",
     ]
 
     # assert we can update the entity, there is no select step
@@ -3174,6 +3178,7 @@ async def test_subentry_reconfigure_add_entity(
         "entity",
         "update_entity",
         "device",
+        "availability",
     ]
 
     # assert we can update the entity, there is no select step
@@ -3272,6 +3277,7 @@ async def test_subentry_reconfigure_update_device_properties(
         "update_entity",
         "delete_entity",
         "device",
+        "availability",
     ]
 
     # assert we can update the device properties
@@ -3310,3 +3316,119 @@ async def test_subentry_reconfigure_update_device_properties(
     assert "hw_version" not in device
     assert device["model"] == "Beer bottle XL"
     assert device["model_id"] == "bn003"
+
+
+@pytest.mark.parametrize(
+    "mqtt_config_subentries_data",
+    [
+        (
+            ConfigSubentryData(
+                data=MOCK_NOTIFY_SUBENTRY_DATA_MULTI,
+                subentry_type="device",
+                title="Mock subentry",
+            ),
+        )
+    ],
+)
+async def test_subentry_reconfigure_availablity(
+    hass: HomeAssistant,
+    mqtt_mock_entry: MqttMockHAClientGenerator,
+) -> None:
+    """Test the subentry ConfigFlow reconfigure and update device properties."""
+    await mqtt_mock_entry()
+    config_entry: MockConfigEntry = hass.config_entries.async_entries(mqtt.DOMAIN)[0]
+    subentry_id: str
+    subentry: ConfigSubentry
+    subentry_id, subentry = next(iter(config_entry.subentries.items()))
+
+    expected_availability = {
+        "availability_topic": "test/availability",
+        "availability_template": "{{ value_json.availability }}",
+        "payload_available": "online",
+        "payload_not_available": "offline",
+    }
+    assert subentry.data.get("availability") == expected_availability
+
+    result = await config_entry.start_subentry_reconfigure_flow(
+        hass, "device", subentry_id
+    )
+    assert result["type"] is FlowResultType.MENU
+    assert result["step_id"] == "summary_menu"
+
+    # assert we can set the availability config
+    result = await hass.config_entries.subentries.async_configure(
+        result["flow_id"],
+        {"next_step_id": "availability"},
+    )
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "availability"
+
+    result = await hass.config_entries.subentries.async_configure(
+        result["flow_id"],
+        user_input={
+            "availability_topic": "test/new_availability#invalid_topic",
+            "payload_available": "1",
+            "payload_not_available": "0",
+        },
+    )
+    assert result["errors"] == {"availability_topic": "invalid_subscribe_topic"}
+
+    result = await hass.config_entries.subentries.async_configure(
+        result["flow_id"],
+        user_input={
+            "availability_topic": "test/new_availability",
+            "payload_available": "1",
+            "payload_not_available": "0",
+        },
+    )
+
+    # finish reconfigure flow
+    result = await hass.config_entries.subentries.async_configure(
+        result["flow_id"],
+        {"next_step_id": "save_changes"},
+    )
+    assert result["type"] is FlowResultType.ABORT
+    assert result["reason"] == "reconfigure_successful"
+
+    # Check the availability was updated
+    expected_availability = {
+        "availability_topic": "test/new_availability",
+        "payload_available": "1",
+        "payload_not_available": "0",
+    }
+    assert subentry.data.get("availability") == expected_availability
+
+    # Assert we can reset the availability config
+    result = await config_entry.start_subentry_reconfigure_flow(
+        hass, "device", subentry_id
+    )
+    assert result["type"] is FlowResultType.MENU
+    assert result["step_id"] == "summary_menu"
+    result = await hass.config_entries.subentries.async_configure(
+        result["flow_id"],
+        {"next_step_id": "availability"},
+    )
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "availability"
+
+    result = await hass.config_entries.subentries.async_configure(
+        result["flow_id"],
+        user_input={
+            "payload_available": "1",
+            "payload_not_available": "0",
+        },
+    )
+
+    # Finish reconfigure flow
+    result = await hass.config_entries.subentries.async_configure(
+        result["flow_id"],
+        {"next_step_id": "save_changes"},
+    )
+    assert result["type"] is FlowResultType.ABORT
+    assert result["reason"] == "reconfigure_successful"
+
+    # Check the availability was updated
+    assert subentry.data.get("availability") == {
+        "payload_available": "1",
+        "payload_not_available": "0",
+    }

tests/components/mqtt/test_mixins.py

@@ -501,6 +501,20 @@ async def test_loading_subentries(
         assert entity_entry_entity_id == entity_id
         state = hass.states.get(entity_id)
         assert state is not None
+        assert (
+            state.attributes.get("entity_picture") == f"https://example.com/{object_id}"
+        )
+        # Availability was configured, so entities are unavailable
+        assert state.state == "unavailable"
+
+    # Make entities available
+    async_fire_mqtt_message(hass, "test/availability", '{"availability": "online"}')
+    for component in mqtt_config_subentries_data[0]["data"]["components"].values():
+        platform = component["platform"]
+        entity_id = f"{platform}.{slugify(device.name)}_{slugify(component['name'])}"
+        state = hass.states.get(entity_id)
+        assert state is not None
+        assert state.state == "unknown"
 
 
 @pytest.mark.parametrize(