From 19040becd37f1569cbb4fd66b5a0f3de817290bb Mon Sep 17 00:00:00 2001
From: Jan Bouwhuis <jbouwh@users.noreply.github.com>
Date: Fri, 24 Nov 2023 21:54:53 +0100
Subject: [PATCH] Fix hassio mqtt discovery CI test (#104463)

* Fix hassio mqtt discovery CI test

* Avoid mqtt set up before mocking the flow

* Fix mock
---
 tests/common.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tests/common.py b/tests/common.py
index 06cdee06ec9..30ea779295c 100644
--- a/tests/common.py
+++ b/tests/common.py
@@ -1304,11 +1304,12 @@ async def get_system_health_info(hass: HomeAssistant, domain: str) -> dict[str,
 @contextmanager
 def mock_config_flow(domain: str, config_flow: type[ConfigFlow]) -> None:
     """Mock a config flow handler."""
-    assert domain not in config_entries.HANDLERS
+    handler = config_entries.HANDLERS.get(domain)
     config_entries.HANDLERS[domain] = config_flow
     _LOGGER.info("Adding mock config flow: %s", domain)
     yield
-    config_entries.HANDLERS.pop(domain)
+    if handler:
+        config_entries.HANDLERS[domain] = handler
 
 
 def mock_integration(