Convert mqtt entity discovery to use callbacks (#118200)

homeassistant/components/mqtt/alarm_control_panel.py

@@ -127,7 +127,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up MQTT alarm control panel through YAML and through MQTT discovery."""
-    await async_setup_entity_entry_helper(
+    async_setup_entity_entry_helper(
         hass,
         config_entry,
         MqttAlarm,

homeassistant/components/mqtt/binary_sensor.py

@@ -71,7 +71,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up MQTT binary sensor through YAML and through MQTT discovery."""
-    await async_setup_entity_entry_helper(
+    async_setup_entity_entry_helper(
         hass,
         config_entry,
         MqttBinarySensor,

homeassistant/components/mqtt/button.py

@@ -44,7 +44,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up MQTT button through YAML and through MQTT discovery."""
-    await async_setup_entity_entry_helper(
+    async_setup_entity_entry_helper(
         hass,
         config_entry,
         MqttButton,

homeassistant/components/mqtt/camera.py

@@ -61,7 +61,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up MQTT camera through YAML and through MQTT discovery."""
-    await async_setup_entity_entry_helper(
+    async_setup_entity_entry_helper(
         hass,
         config_entry,
         MqttCamera,

homeassistant/components/mqtt/climate.py

@@ -379,7 +379,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up MQTT climate through YAML and through MQTT discovery."""
-    await async_setup_entity_entry_helper(
+    async_setup_entity_entry_helper(
         hass,
         config_entry,
         MqttClimate,

homeassistant/components/mqtt/cover.py

@@ -220,7 +220,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up MQTT cover through YAML and through MQTT discovery."""
-    await async_setup_entity_entry_helper(
+    async_setup_entity_entry_helper(
         hass,
         config_entry,
         MqttCover,

homeassistant/components/mqtt/device_automation.py

@@ -29,7 +29,7 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> N
     """Set up MQTT device automation dynamically through MQTT discovery."""
 
     setup = functools.partial(_async_setup_automation, hass, config_entry=config_entry)
-    await async_setup_non_entity_entry_helper(
+    async_setup_non_entity_entry_helper(
         hass, "device_automation", setup, DISCOVERY_SCHEMA
     )
 

homeassistant/components/mqtt/device_tracker.py

@@ -83,7 +83,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up MQTT event through YAML and through MQTT discovery."""
-    await async_setup_entity_entry_helper(
+    async_setup_entity_entry_helper(
         hass,
         config_entry,
         MqttDeviceTracker,

homeassistant/components/mqtt/event.py

@@ -74,7 +74,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up MQTT event through YAML and through MQTT discovery."""
-    await async_setup_entity_entry_helper(
+    async_setup_entity_entry_helper(
         hass,
         config_entry,
         MqttEvent,

homeassistant/components/mqtt/fan.py

@@ -191,7 +191,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up MQTT fan through YAML and through MQTT discovery."""
-    await async_setup_entity_entry_helper(
+    async_setup_entity_entry_helper(
         hass,
         config_entry,
         MqttFan,

homeassistant/components/mqtt/humidifier.py

@@ -184,7 +184,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up MQTT humidifier through YAML and through MQTT discovery."""
-    await async_setup_entity_entry_helper(
+    async_setup_entity_entry_helper(
         hass,
         config_entry,
         MqttHumidifier,

homeassistant/components/mqtt/image.py

@@ -83,7 +83,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up MQTT image through YAML and through MQTT discovery."""
-    await async_setup_entity_entry_helper(
+    async_setup_entity_entry_helper(
         hass,
         config_entry,
         MqttImage,

homeassistant/components/mqtt/lawn_mower.py

@@ -81,7 +81,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up MQTT lawn mower through YAML and through MQTT discovery."""
-    await async_setup_entity_entry_helper(
+    async_setup_entity_entry_helper(
         hass,
         config_entry,
         MqttLawnMower,

homeassistant/components/mqtt/light/__init__.py

@@ -70,7 +70,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up MQTT lights through YAML and through MQTT discovery."""
-    await async_setup_entity_entry_helper(
+    async_setup_entity_entry_helper(
         hass,
         config_entry,
         None,

homeassistant/components/mqtt/lock.py

@@ -117,7 +117,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up MQTT lock through YAML and through MQTT discovery."""
-    await async_setup_entity_entry_helper(
+    async_setup_entity_entry_helper(
         hass,
         config_entry,
         MqttLock,

homeassistant/components/mqtt/mixins.py

@@ -4,7 +4,6 @@ from __future__ import annotations
 
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Coroutine
-import functools
 from functools import partial
 import logging
 from typing import TYPE_CHECKING, Any, Protocol, cast, final
@@ -169,17 +168,20 @@ def async_handle_schema_error(
     )
 
 
-async def _async_discover(
+def _handle_discovery_failure(
     hass: HomeAssistant,
-    domain: str,
-    setup: Callable[[MQTTDiscoveryPayload], None] | None,
-    async_setup: Callable[[MQTTDiscoveryPayload], Coroutine[Any, Any, None]] | None,
     discovery_payload: MQTTDiscoveryPayload,
 ) -> None:
-    """Discover and add an MQTT entity, automation or tag.
+    """Handle discovery failure."""
+    discovery_hash = discovery_payload.discovery_data[ATTR_DISCOVERY_HASH]
+    clear_discovery_hash(hass, discovery_hash)
+    async_dispatcher_send(hass, MQTT_DISCOVERY_DONE.format(*discovery_hash), None)
 
-    setup is to be run in the event loop when there is nothing to be awaited.
-    """
+
+def _verify_mqtt_config_entry_enabled_for_discovery(
+    hass: HomeAssistant, domain: str, discovery_payload: MQTTDiscoveryPayload
+) -> bool:
+    """Verify MQTT config entry is enabled or log warning."""
     if not mqtt_config_entry_enabled(hass):
         _LOGGER.warning(
             (
@@ -189,23 +191,8 @@ async def _async_discover(
             domain,
             discovery_payload,
         )
-        return
-    discovery_data = discovery_payload.discovery_data
-    try:
-        if setup is not None:
-            setup(discovery_payload)
-        elif async_setup is not None:
-            await async_setup(discovery_payload)
-    except vol.Invalid as err:
-        discovery_hash = discovery_data[ATTR_DISCOVERY_HASH]
-        clear_discovery_hash(hass, discovery_hash)
-        async_dispatcher_send(hass, MQTT_DISCOVERY_DONE.format(*discovery_hash), None)
-        async_handle_schema_error(discovery_payload, err)
-    except Exception:
-        discovery_hash = discovery_data[ATTR_DISCOVERY_HASH]
-        clear_discovery_hash(hass, discovery_hash)
-        async_dispatcher_send(hass, MQTT_DISCOVERY_DONE.format(*discovery_hash), None)
-        raise
+        return False
+    return True
 
 
 class _SetupNonEntityHelperCallbackProtocol(Protocol):  # pragma: no cover
@@ -216,7 +203,8 @@ class _SetupNonEntityHelperCallbackProtocol(Protocol):  # pragma: no cover
     ) -> None: ...
 
 
-async def async_setup_non_entity_entry_helper(
+@callback
+def async_setup_non_entity_entry_helper(
     hass: HomeAssistant,
     domain: str,
     async_setup: _SetupNonEntityHelperCallbackProtocol,
@@ -225,25 +213,35 @@ async def async_setup_non_entity_entry_helper(
     """Set up automation or tag creation dynamically through MQTT discovery."""
     mqtt_data = hass.data[DATA_MQTT]
 
-    async def async_setup_from_discovery(
+    async def _async_setup_non_entity_entry_from_discovery(
         discovery_payload: MQTTDiscoveryPayload,
     ) -> None:
         """Set up an MQTT entity, automation or tag from discovery."""
-        config: ConfigType = discovery_schema(discovery_payload)
-        await async_setup(config, discovery_data=discovery_payload.discovery_data)
+        if not _verify_mqtt_config_entry_enabled_for_discovery(
+            hass, domain, discovery_payload
+        ):
+            return
+        try:
+            config: ConfigType = discovery_schema(discovery_payload)
+            await async_setup(config, discovery_data=discovery_payload.discovery_data)
+        except vol.Invalid as err:
+            _handle_discovery_failure(hass, discovery_payload)
+            async_handle_schema_error(discovery_payload, err)
+        except Exception:
+            _handle_discovery_failure(hass, discovery_payload)
+            raise
 
     mqtt_data.reload_dispatchers.append(
         async_dispatcher_connect(
             hass,
             MQTT_DISCOVERY_NEW.format(domain, "mqtt"),
-            functools.partial(
-                _async_discover, hass, domain, None, async_setup_from_discovery
-            ),
+            _async_setup_non_entity_entry_from_discovery,
         )
     )
 
 
-async def async_setup_entity_entry_helper(
+@callback
+def async_setup_entity_entry_helper(
     hass: HomeAssistant,
     entry: ConfigEntry,
     entity_class: type[MqttEntity] | None,
@@ -257,27 +255,36 @@ async def async_setup_entity_entry_helper(
     mqtt_data = hass.data[DATA_MQTT]
 
     @callback
-    def async_setup_from_discovery(
+    def _async_setup_entity_entry_from_discovery(
         discovery_payload: MQTTDiscoveryPayload,
     ) -> None:
         """Set up an MQTT entity from discovery."""
         nonlocal entity_class
-        config: DiscoveryInfoType = discovery_schema(discovery_payload)
-        if schema_class_mapping is not None:
-            entity_class = schema_class_mapping[config[CONF_SCHEMA]]
-        if TYPE_CHECKING:
-            assert entity_class is not None
-        async_add_entities(
-            [entity_class(hass, config, entry, discovery_payload.discovery_data)]
-        )
+        if not _verify_mqtt_config_entry_enabled_for_discovery(
+            hass, domain, discovery_payload
+        ):
+            return
+        try:
+            config: DiscoveryInfoType = discovery_schema(discovery_payload)
+            if schema_class_mapping is not None:
+                entity_class = schema_class_mapping[config[CONF_SCHEMA]]
+            if TYPE_CHECKING:
+                assert entity_class is not None
+            async_add_entities(
+                [entity_class(hass, config, entry, discovery_payload.discovery_data)]
+            )
+        except vol.Invalid as err:
+            _handle_discovery_failure(hass, discovery_payload)
+            async_handle_schema_error(discovery_payload, err)
+        except Exception:
+            _handle_discovery_failure(hass, discovery_payload)
+            raise
 
     mqtt_data.reload_dispatchers.append(
         async_dispatcher_connect(
             hass,
             MQTT_DISCOVERY_NEW.format(domain, "mqtt"),
-            functools.partial(
-                _async_discover, hass, domain, async_setup_from_discovery, None
-            ),
+            _async_setup_entity_entry_from_discovery,
         )
     )
 

homeassistant/components/mqtt/notify.py

@@ -40,7 +40,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up MQTT notify through YAML and through MQTT discovery."""
-    await async_setup_entity_entry_helper(
+    async_setup_entity_entry_helper(
         hass,
         config_entry,
         MqttNotify,

homeassistant/components/mqtt/number.py

@@ -110,7 +110,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up MQTT number through YAML and through MQTT discovery."""
-    await async_setup_entity_entry_helper(
+    async_setup_entity_entry_helper(
         hass,
         config_entry,
         MqttNumber,

homeassistant/components/mqtt/scene.py

@@ -44,7 +44,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up MQTT scene through YAML and through MQTT discovery."""
-    await async_setup_entity_entry_helper(
+    async_setup_entity_entry_helper(
         hass,
         config_entry,
         MqttScene,

homeassistant/components/mqtt/select.py

@@ -61,7 +61,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up MQTT select through YAML and through MQTT discovery."""
-    await async_setup_entity_entry_helper(
+    async_setup_entity_entry_helper(
         hass,
         config_entry,
         MqttSelect,

homeassistant/components/mqtt/sensor.py

@@ -110,7 +110,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up MQTT sensor through YAML and through MQTT discovery."""
-    await async_setup_entity_entry_helper(
+    async_setup_entity_entry_helper(
         hass,
         config_entry,
         MqttSensor,

homeassistant/components/mqtt/siren.py

@@ -114,7 +114,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up MQTT siren through YAML and through MQTT discovery."""
-    await async_setup_entity_entry_helper(
+    async_setup_entity_entry_helper(
         hass,
         config_entry,
         MqttSiren,

homeassistant/components/mqtt/switch.py

@@ -60,7 +60,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up MQTT switch through YAML and through MQTT discovery."""
-    await async_setup_entity_entry_helper(
+    async_setup_entity_entry_helper(
         hass,
         config_entry,
         MqttSwitch,

homeassistant/components/mqtt/tag.py

@@ -57,7 +57,7 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> N
     """Set up MQTT tag scanner dynamically through MQTT discovery."""
 
     setup = functools.partial(_async_setup_tag, hass, config_entry=config_entry)
-    await async_setup_non_entity_entry_helper(hass, TAG, setup, DISCOVERY_SCHEMA)
+    async_setup_non_entity_entry_helper(hass, TAG, setup, DISCOVERY_SCHEMA)
 
 
 async def _async_setup_tag(

homeassistant/components/mqtt/text.py

@@ -96,7 +96,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up MQTT text through YAML and through MQTT discovery."""
-    await async_setup_entity_entry_helper(
+    async_setup_entity_entry_helper(
         hass,
         config_entry,
         MqttTextEntity,

homeassistant/components/mqtt/update.py

@@ -80,7 +80,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up MQTT update entity through YAML and through MQTT discovery."""
-    await async_setup_entity_entry_helper(
+    async_setup_entity_entry_helper(
         hass,
         config_entry,
         MqttUpdate,

homeassistant/components/mqtt/vacuum.py

@@ -236,7 +236,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up MQTT vacuum through YAML and through MQTT discovery."""
-    await async_setup_entity_entry_helper(
+    async_setup_entity_entry_helper(
         hass,
         config_entry,
         MqttStateVacuum,

homeassistant/components/mqtt/valve.py

@@ -140,7 +140,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up MQTT valve through YAML and through MQTT discovery."""
-    await async_setup_entity_entry_helper(
+    async_setup_entity_entry_helper(
         hass,
         config_entry,
         MqttValve,

homeassistant/components/mqtt/water_heater.py

@@ -167,7 +167,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up MQTT water heater device through YAML and through MQTT discovery."""
-    await async_setup_entity_entry_helper(
+    async_setup_entity_entry_helper(
         hass,
         config_entry,
         MqttWaterHeater,