Clean up MQTT platform entry setup at discovery (#72371)

* Setup MQTT discovery with entry setup

* Wait for entry setup in test

* flake

homeassistant/components/mqtt/__init__.py

@@ -236,9 +236,7 @@ async def _async_config_entry_updated(hass: HomeAssistant, entry: ConfigEntry) -
         await _async_setup_discovery(hass, mqtt_client.conf, entry)
 
 
-async def async_setup_entry(  # noqa: C901
+async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
-    hass: HomeAssistant, entry: ConfigEntry
-) -> bool:
     """Load a config entry."""
     # Merge basic configuration, and add missing defaults for basic options
     _merge_basic_config(hass, entry, hass.data.get(DATA_MQTT_CONFIG, {}))
@@ -386,25 +384,33 @@ async def async_setup_entry(  # noqa: C901
         hass.data[DATA_MQTT_UPDATED_CONFIG] = config_yaml.get(DOMAIN, {})
         async_dispatcher_send(hass, MQTT_RELOADED)
 
-    async def async_forward_entry_setup():
+    async def async_forward_entry_setup_and_setup_discovery(config_entry):
-        """Forward the config entry setup to the platforms."""
+        """Forward the config entry setup to the platforms and set up discovery."""
-        async with hass.data[DATA_CONFIG_ENTRY_LOCK]:
+        # Local import to avoid circular dependencies
-            for component in PLATFORMS:
+        # pylint: disable-next=import-outside-toplevel
-                config_entries_key = f"{component}.mqtt"
+        from . import device_automation, tag
-                if config_entries_key not in hass.data[CONFIG_ENTRY_IS_SETUP]:
+
-                    hass.data[CONFIG_ENTRY_IS_SETUP].add(config_entries_key)
+        await asyncio.gather(
-                    await hass.config_entries.async_forward_entry_setup(
+            *(
-                        entry, component
+                [
-                    )
+                    device_automation.async_setup_entry(hass, config_entry),
+                    tag.async_setup_entry(hass, config_entry),
+                ]
+                + [
+                    hass.config_entries.async_forward_entry_setup(entry, component)
+                    for component in PLATFORMS
+                ]
+            )
+        )
+        # Setup discovery
+        if conf.get(CONF_DISCOVERY):
+            await _async_setup_discovery(hass, conf, entry)
         # Setup reload service after all platforms have loaded
         entry.async_on_unload(
             hass.bus.async_listen("event_mqtt_reloaded", _async_reload_platforms)
         )
 
-    hass.async_create_task(async_forward_entry_setup())
+    hass.async_create_task(async_forward_entry_setup_and_setup_discovery(entry))
-
-    if conf.get(CONF_DISCOVERY):
-        await _async_setup_discovery(hass, conf, entry)
 
     if DATA_MQTT_RELOAD_NEEDED in hass.data:
         hass.data.pop(DATA_MQTT_RELOAD_NEEDED)

homeassistant/components/mqtt/discovery.py

@@ -27,8 +27,6 @@ from .const import (
     ATTR_DISCOVERY_TOPIC,
     CONF_AVAILABILITY,
     CONF_TOPIC,
-    CONFIG_ENTRY_IS_SETUP,
-    DATA_CONFIG_ENTRY_LOCK,
     DOMAIN,
 )
 
@@ -227,28 +225,6 @@ async def async_start(  # noqa: C901
             # Add component
             _LOGGER.info("Found new component: %s %s", component, discovery_id)
             hass.data[ALREADY_DISCOVERED][discovery_hash] = None
-
-            config_entries_key = f"{component}.mqtt"
-            async with hass.data[DATA_CONFIG_ENTRY_LOCK]:
-                if config_entries_key not in hass.data[CONFIG_ENTRY_IS_SETUP]:
-                    if component == "device_automation":
-                        # Local import to avoid circular dependencies
-                        # pylint: disable-next=import-outside-toplevel
-                        from . import device_automation
-
-                        await device_automation.async_setup_entry(hass, config_entry)
-                    elif component == "tag":
-                        # Local import to avoid circular dependencies
-                        # pylint: disable-next=import-outside-toplevel
-                        from . import tag
-
-                        await tag.async_setup_entry(hass, config_entry)
-                    else:
-                        await hass.config_entries.async_forward_entry_setup(
-                            config_entry, component
-                        )
-                    hass.data[CONFIG_ENTRY_IS_SETUP].add(config_entries_key)
-
             async_dispatcher_send(
                 hass, MQTT_DISCOVERY_NEW.format(component, "mqtt"), payload
             )

tests/components/mqtt/test_init.py

@@ -1780,13 +1780,12 @@ async def test_setup_entry_with_config_override(
     # mqtt present in yaml config
     assert await async_setup_component(hass, mqtt.DOMAIN, {})
     await hass.async_block_till_done()
-    await mqtt_mock_entry_with_yaml_config()
 
     # User sets up a config entry
     entry = MockConfigEntry(domain=mqtt.DOMAIN, data={mqtt.CONF_BROKER: "test-broker"})
     entry.add_to_hass(hass)
-    with patch("homeassistant.components.mqtt.PLATFORMS", []):
+    assert await hass.config_entries.async_setup(entry.entry_id)
-        assert await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
 
     # Discover a device to verify the entry was setup correctly
     async_fire_mqtt_message(hass, "homeassistant/sensor/bla/config", data)