Allow Just-in-Time platform setup for mqtt (#112720)

* Allow Just-in-Time platform setup for mqtt

* Only forward the setup of new platforms

* Fix new  platforms being setup at reload + test

* Revert not related changes

* Remove unused partial

* Address comments, only import plaforms if needed

* Apply suggestions from code review

* Add multipl platform discovery test

* Improve test

* Use a lock per platform

tests/components/mqtt/test_init.py

@@ -2,6 +2,7 @@
 
 import asyncio
 from collections.abc import Generator
+from copy import deepcopy
 from datetime import datetime, timedelta
 from functools import partial
 import json
@@ -3546,7 +3547,6 @@ async def test_subscribe_connection_status(
     assert mqtt_connected_calls_async[1] is False
 
 
-@patch("homeassistant.components.mqtt.PLATFORMS", [Platform.LIGHT])
 async def test_unload_config_entry(
     hass: HomeAssistant,
     mqtt_mock: MqttMockHAClient,
@@ -3563,6 +3563,7 @@ async def test_unload_config_entry(
     # Publish just before unloading to test await cleanup
     mqtt_client_mock.reset_mock()
     mqtt.publish(hass, "just_in_time", "published", qos=0, retain=False)
+    await hass.async_block_till_done()
 
     assert await hass.config_entries.async_unload(mqtt_config_entry.entry_id)
     new_mqtt_config_entry = mqtt_config_entry
@@ -4046,3 +4047,127 @@ async def test_reload_with_empty_config(
         await hass.async_block_till_done()
 
     assert hass.states.get("sensor.test") is None
+
+
+@pytest.mark.parametrize(
+    "hass_config",
+    [
+        {
+            "mqtt": [
+                {
+                    "sensor": {
+                        "name": "test",
+                        "state_topic": "test-topic",
+                    }
+                },
+            ]
+        }
+    ],
+)
+async def test_reload_with_new_platform_config(
+    hass: HomeAssistant,
+    mqtt_mock_entry: MqttMockHAClientGenerator,
+) -> None:
+    """Test reloading yaml with new platform config."""
+    await mqtt_mock_entry()
+    assert hass.states.get("sensor.test") is not None
+    assert hass.states.get("binary_sensor.test") is None
+
+    new_config = {
+        "mqtt": [
+            {
+                "sensor": {
+                    "name": "test",
+                    "state_topic": "test-topic1",
+                },
+                "binary_sensor": {
+                    "name": "test",
+                    "state_topic": "test-topic2",
+                },
+            },
+        ]
+    }
+
+    # Reload with an new platform config and assert again
+    with patch("homeassistant.config.load_yaml_config_file", return_value=new_config):
+        await hass.services.async_call(
+            "mqtt",
+            SERVICE_RELOAD,
+            {},
+            blocking=True,
+        )
+        await hass.async_block_till_done()
+
+    assert hass.states.get("sensor.test") is not None
+    assert hass.states.get("binary_sensor.test") is not None
+
+
+async def test_multi_platform_discovery(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    mqtt_mock_entry: MqttMockHAClientGenerator,
+) -> None:
+    """Test setting up multiple platforms simultaneous."""
+    await mqtt_mock_entry()
+    entity_configs = {
+        "alarm_control_panel": {
+            "name": "test",
+            "state_topic": "alarm/state",
+            "command_topic": "alarm/command",
+        },
+        "button": {"name": "test", "command_topic": "test-topic"},
+        "camera": {"name": "test", "topic": "test_topic"},
+        "cover": {"name": "test", "state_topic": "test-topic"},
+        "device_tracker": {
+            "name": "test",
+            "state_topic": "test-topic",
+        },
+        "fan": {
+            "name": "test",
+            "state_topic": "state-topic",
+            "command_topic": "command-topic",
+        },
+        "sensor": {"name": "test", "state_topic": "test-topic"},
+        "switch": {"name": "test", "command_topic": "test-topic"},
+        "select": {
+            "name": "test",
+            "command_topic": "test-topic",
+            "options": ["milk", "beer"],
+        },
+    }
+    non_entity_configs = {
+        "tag": {
+            "device": {"identifiers": ["tag_0AFFD2"]},
+            "topic": "foobar/tag_scanned",
+        },
+        "device_automation": {
+            "automation_type": "trigger",
+            "device": {"identifiers": ["device_automation_0AFFD2"]},
+            "payload": "short_press",
+            "topic": "foobar/triggers/button1",
+            "type": "button_short_press",
+            "subtype": "button_1",
+        },
+    }
+    for platform, config in entity_configs.items():
+        for set_number in range(0, 2):
+            set_config = deepcopy(config)
+            set_config["name"] = f"test_{set_number}"
+            topic = f"homeassistant/{platform}/bla_{set_number}/config"
+            async_fire_mqtt_message(hass, topic, json.dumps(set_config))
+    for platform, config in non_entity_configs.items():
+        topic = f"homeassistant/{platform}/bla/config"
+        async_fire_mqtt_message(hass, topic, json.dumps(config))
+    await hass.async_block_till_done()
+    for set_number in range(0, 2):
+        for platform in entity_configs:
+            entity_id = f"{platform}.test_{set_number}"
+            state = hass.states.get(entity_id)
+            assert state is not None
+    for platform in non_entity_configs:
+        assert (
+            device_registry.async_get_device(
+                identifiers={("mqtt", f"{platform}_0AFFD2")}
+            )
+            is not None
+        )