Fix setup and tear down issues for mqtt discovery and config flow tests (#120333)

* Fix setup and tear down issues for mqtt discovery and config flow tests

* Use async callback

tests/components/mqtt/conftest.py

@@ -6,7 +6,17 @@ from unittest.mock import patch
 import pytest
 from typing_extensions import Generator
 
-from tests.components.light.conftest import mock_light_profiles  # noqa: F401
+from homeassistant.components import mqtt
+
+ENTRY_DEFAULT_BIRTH_MESSAGE = {
+    mqtt.CONF_BROKER: "mock-broker",
+    mqtt.CONF_BIRTH_MESSAGE: {
+        mqtt.ATTR_TOPIC: "homeassistant/status",
+        mqtt.ATTR_PAYLOAD: "online",
+        mqtt.ATTR_QOS: 0,
+        mqtt.ATTR_RETAIN: False,
+    },
+}
 
 
 @pytest.fixture(autouse=True)

tests/components/mqtt/test_config_flow.py

@@ -1104,7 +1104,6 @@ async def test_skipping_advanced_options(
 )
 async def test_step_reauth(
     hass: HomeAssistant,
-    mqtt_mock_entry: MqttMockHAClientGenerator,
     mqtt_client_mock: MqttMockPahoClient,
     mock_try_connection: MagicMock,
     mock_reload_after_entry_update: MagicMock,
@@ -1115,12 +1114,9 @@ async def test_step_reauth(
     """Test that the reauth step works."""
 
     # Prepare the config entry
-    config_entry = hass.config_entries.async_entries(mqtt.DOMAIN)[0]
+    config_entry = MockConfigEntry(domain=mqtt.DOMAIN, data=test_input)
-    hass.config_entries.async_update_entry(
+    config_entry.add_to_hass(hass)
-        config_entry,
+    assert await hass.config_entries.async_setup(config_entry.entry_id)
-        data=test_input,
-    )
-    await mqtt_mock_entry()
 
     # Start reauth flow
     config_entry.async_start_reauth(hass)

tests/components/mqtt/test_discovery.py

@@ -22,7 +22,9 @@ from homeassistant.components.mqtt.discovery import (
     MQTTDiscoveryPayload,
     async_start,
 )
+from homeassistant.components.mqtt.models import ReceiveMessage
 from homeassistant.const import (
+    EVENT_HOMEASSISTANT_STARTED,
     EVENT_STATE_CHANGED,
     STATE_ON,
     STATE_UNAVAILABLE,
@@ -40,6 +42,7 @@ from homeassistant.helpers.service_info.mqtt import MqttServiceInfo
 from homeassistant.setup import async_setup_component
 from homeassistant.util.signal_type import SignalTypeFormat
 
+from .conftest import ENTRY_DEFAULT_BIRTH_MESSAGE
 from .test_common import help_all_subscribe_calls, help_test_unload_config_entry
 
 from tests.common import (
@@ -1454,32 +1457,15 @@ async def test_complex_discovery_topic_prefix(
     ].discovery_already_discovered
 
 
+@patch("homeassistant.components.mqtt.client.DISCOVERY_COOLDOWN", 0.0)
 @patch("homeassistant.components.mqtt.client.INITIAL_SUBSCRIBE_COOLDOWN", 0.0)
 @patch("homeassistant.components.mqtt.client.SUBSCRIBE_COOLDOWN", 0.0)
 @patch("homeassistant.components.mqtt.client.UNSUBSCRIBE_COOLDOWN", 0.0)
 async def test_mqtt_integration_discovery_subscribe_unsubscribe(
     hass: HomeAssistant,
     mqtt_client_mock: MqttMockPahoClient,
-    mqtt_mock_entry: MqttMockHAClientGenerator,
 ) -> None:
     """Check MQTT integration discovery subscribe and unsubscribe."""
-    mqtt_mock = await mqtt_mock_entry()
-    mock_platform(hass, "comp.config_flow", None)
-
-    entry = hass.config_entries.async_entries("mqtt")[0]
-    mqtt_mock().connected = True
-
-    with patch(
-        "homeassistant.components.mqtt.discovery.async_get_mqtt",
-        return_value={"comp": ["comp/discovery/#"]},
-    ):
-        await async_start(hass, "homeassistant", entry)
-        await hass.async_block_till_done()
-        await hass.async_block_till_done()
-        await hass.async_block_till_done()
-
-    assert ("comp/discovery/#", 0) in help_all_subscribe_calls(mqtt_client_mock)
-    assert not mqtt_client_mock.unsubscribe.called
 
     class TestFlow(config_entries.ConfigFlow):
         """Test flow."""
@@ -1488,49 +1474,57 @@ async def test_mqtt_integration_discovery_subscribe_unsubscribe(
             """Test mqtt step."""
             return self.async_abort(reason="already_configured")
 
-    assert not mqtt_client_mock.unsubscribe.called
+    mock_platform(hass, "comp.config_flow", None)
+
+    birth = asyncio.Event()
+
+    @callback
+    def wait_birth(msg: ReceiveMessage) -> None:
+        """Handle birth message."""
+        birth.set()
 
     wait_unsub = asyncio.Event()
 
+    @callback
     def _mock_unsubscribe(topics: list[str]) -> tuple[int, int]:
         wait_unsub.set()
         return (0, 0)
 
+    entry = MockConfigEntry(domain=mqtt.DOMAIN, data=ENTRY_DEFAULT_BIRTH_MESSAGE)
+    entry.add_to_hass(hass)
     with (
+        patch(
+            "homeassistant.components.mqtt.discovery.async_get_mqtt",
+            return_value={"comp": ["comp/discovery/#"]},
+        ),
         mock_config_flow("comp", TestFlow),
         patch.object(mqtt_client_mock, "unsubscribe", side_effect=_mock_unsubscribe),
     ):
+        assert await hass.config_entries.async_setup(entry.entry_id)
+        await mqtt.async_subscribe(hass, "homeassistant/status", wait_birth)
+        hass.bus.async_fire(EVENT_HOMEASSISTANT_STARTED)
+        await birth.wait()
+
+        assert ("comp/discovery/#", 0) in help_all_subscribe_calls(mqtt_client_mock)
+        assert not mqtt_client_mock.unsubscribe.called
+        mqtt_client_mock.reset_mock()
+
+        await hass.async_block_till_done(wait_background_tasks=True)
         async_fire_mqtt_message(hass, "comp/discovery/bla/config", "")
         await wait_unsub.wait()
         mqtt_client_mock.unsubscribe.assert_called_once_with(["comp/discovery/#"])
+        await hass.async_block_till_done(wait_background_tasks=True)
 
 
+@patch("homeassistant.components.mqtt.client.DISCOVERY_COOLDOWN", 0.0)
 @patch("homeassistant.components.mqtt.client.INITIAL_SUBSCRIBE_COOLDOWN", 0.0)
 @patch("homeassistant.components.mqtt.client.SUBSCRIBE_COOLDOWN", 0.0)
 @patch("homeassistant.components.mqtt.client.UNSUBSCRIBE_COOLDOWN", 0.0)
 async def test_mqtt_discovery_unsubscribe_once(
     hass: HomeAssistant,
     mqtt_client_mock: MqttMockPahoClient,
-    mqtt_mock_entry: MqttMockHAClientGenerator,
 ) -> None:
     """Check MQTT integration discovery unsubscribe once."""
-    mqtt_mock = await mqtt_mock_entry()
-    mock_platform(hass, "comp.config_flow", None)
-
-    entry = hass.config_entries.async_entries("mqtt")[0]
-    mqtt_mock().connected = True
-
-    with patch(
-        "homeassistant.components.mqtt.discovery.async_get_mqtt",
-        return_value={"comp": ["comp/discovery/#"]},
-    ):
-        await async_start(hass, "homeassistant", entry)
-        await hass.async_block_till_done()
-        await hass.async_block_till_done()
-        await hass.async_block_till_done()
-
-    assert ("comp/discovery/#", 0) in help_all_subscribe_calls(mqtt_client_mock)
-    assert not mqtt_client_mock.unsubscribe.called
 
     class TestFlow(config_entries.ConfigFlow):
         """Test flow."""
@@ -1540,13 +1534,49 @@ async def test_mqtt_discovery_unsubscribe_once(
             await asyncio.sleep(0.1)
             return self.async_abort(reason="already_configured")
 
-    with mock_config_flow("comp", TestFlow):
+    mock_platform(hass, "comp.config_flow", None)
+
+    birth = asyncio.Event()
+
+    @callback
+    def wait_birth(msg: ReceiveMessage) -> None:
+        """Handle birth message."""
+        birth.set()
+
+    wait_unsub = asyncio.Event()
+
+    @callback
+    def _mock_unsubscribe(topics: list[str]) -> tuple[int, int]:
+        wait_unsub.set()
+        return (0, 0)
+
+    entry = MockConfigEntry(domain=mqtt.DOMAIN, data=ENTRY_DEFAULT_BIRTH_MESSAGE)
+    entry.add_to_hass(hass)
+
+    with (
+        patch(
+            "homeassistant.components.mqtt.discovery.async_get_mqtt",
+            return_value={"comp": ["comp/discovery/#"]},
+        ),
+        mock_config_flow("comp", TestFlow),
+        patch.object(mqtt_client_mock, "unsubscribe", side_effect=_mock_unsubscribe),
+    ):
+        assert await hass.config_entries.async_setup(entry.entry_id)
+        await mqtt.async_subscribe(hass, "homeassistant/status", wait_birth)
+        hass.bus.async_fire(EVENT_HOMEASSISTANT_STARTED)
+        await birth.wait()
+
+        assert ("comp/discovery/#", 0) in help_all_subscribe_calls(mqtt_client_mock)
+        assert not mqtt_client_mock.unsubscribe.called
+
+        await hass.async_block_till_done(wait_background_tasks=True)
         async_fire_mqtt_message(hass, "comp/discovery/bla/config", "")
         async_fire_mqtt_message(hass, "comp/discovery/bla/config", "")
-        await asyncio.sleep(0.1)
+        await wait_unsub.wait()
-        await hass.async_block_till_done()
+        await asyncio.sleep(0.2)
-        await hass.async_block_till_done()
+        await hass.async_block_till_done(wait_background_tasks=True)
         mqtt_client_mock.unsubscribe.assert_called_once_with(["comp/discovery/#"])
+        await hass.async_block_till_done(wait_background_tasks=True)
 
 
 async def test_clear_config_topic_disabled_entity(