Pop the mocked config flow, restore the original with mock_config_flow (#107567)

Pop the mocked config flow, restore the original if it existed
This commit is contained in:
Jan Bouwhuis 2024-01-08 19:01:43 +01:00 committed by GitHub
parent 0d946c62dc
commit 20610645fb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -1321,8 +1321,13 @@ async def get_system_health_info(hass: HomeAssistant, domain: str) -> dict[str,
@contextmanager @contextmanager
def mock_config_flow(domain: str, config_flow: type[ConfigFlow]) -> None: def mock_config_flow(domain: str, config_flow: type[ConfigFlow]) -> None:
"""Mock a config flow handler.""" """Mock a config flow handler."""
with patch.dict(config_entries.HANDLERS, {domain: config_flow}): original_handler = config_entries.HANDLERS.get(domain)
yield 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( def mock_integration(