From 20610645fb82e7021c632e57f9a5cc8002ab64bb Mon Sep 17 00:00:00 2001
From: Jan Bouwhuis <jbouwh@users.noreply.github.com>
Date: Mon, 8 Jan 2024 19:01:43 +0100
Subject: [PATCH] Pop the mocked config flow, restore the original with
 mock_config_flow (#107567)

Pop the mocked config flow, restore the original if it existed
---
 tests/common.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/tests/common.py b/tests/common.py
index 4e68bcf4357..02c7150588d 100644
--- a/tests/common.py
+++ b/tests/common.py
@@ -1321,8 +1321,13 @@ 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."""
-    with patch.dict(config_entries.HANDLERS, {domain: config_flow}):
-        yield
+    original_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 original_handler:
+        config_entries.HANDLERS[domain] = original_handler
 
 
 def mock_integration(