From 35900cd579d4b8ad98f5d556840722d6704556b4 Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Sat, 11 May 2024 20:11:42 +0200
Subject: [PATCH] Use mock_config_flow helper in bootstrap tests (#117240)

---
 tests/test_bootstrap.py | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/tests/test_bootstrap.py b/tests/test_bootstrap.py
index 3d2735d9c1c..6e3ec7066e6 100644
--- a/tests/test_bootstrap.py
+++ b/tests/test_bootstrap.py
@@ -27,6 +27,7 @@ from .common import (
     MockModule,
     MockPlatform,
     get_test_config_dir,
+    mock_config_flow,
     mock_integration,
     mock_platform,
 )
@@ -1146,7 +1147,6 @@ async def test_bootstrap_empty_integrations(
 @pytest.fixture(name="mock_mqtt_config_flow")
 def mock_mqtt_config_flow_fixture() -> Generator[None, None, None]:
     """Mock MQTT config flow."""
-    original_mqtt = HANDLERS.get("mqtt")
 
     @HANDLERS.register("mqtt")
     class MockConfigFlow:
@@ -1155,11 +1155,8 @@ def mock_mqtt_config_flow_fixture() -> Generator[None, None, None]:
         VERSION = 1
         MINOR_VERSION = 1
 
-    yield
-    if original_mqtt:
-        HANDLERS["mqtt"] = original_mqtt
-    else:
-        HANDLERS.pop("mqtt")
+    with mock_config_flow("mqtt", MockConfigFlow):
+        yield
 
 
 @pytest.mark.parametrize("integration", ["mqtt_eventstream", "mqtt_statestream"])