Only mock config_entries.HANDLERS for the current test in mock_config…

…_flow (home-assistant#107357)
TNTLarsn · Jan 6, 2024 · d6aaaf1 · d6aaaf1
1 parent 65985c4
commit d6aaaf1
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/tests/common.py b/tests/common.py
@@ -1316,12 +1316,8 @@ 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."""
-    handler = config_entries.HANDLERS.get(domain)
-    config_entries.HANDLERS[domain] = config_flow
-    _LOGGER.info("Adding mock config flow: %s", domain)
-    yield
-    if handler:
-        config_entries.HANDLERS[domain] = handler
+    with patch.dict(config_entries.HANDLERS, {domain: config_flow}):
+        yield
 
 
 def mock_integration(