diff --git a/conftest.py b/conftest.py index ebb9dae1457..966086811f1 100644 --- a/conftest.py +++ b/conftest.py @@ -18,6 +18,13 @@ # Make all fixtures available from distributed.utils_test import * # noqa +from distributed.utils_test import clean_config + + +@pytest.fixture(autouse=True) +def ensure_config_reset(): + with clean_config(): + yield def pytest_addoption(parser): diff --git a/distributed/utils_test.py b/distributed/utils_test.py index f6625e74180..19be64b42c5 100644 --- a/distributed/utils_test.py +++ b/distributed/utils_test.py @@ -19,6 +19,7 @@ import tempfile import threading import uuid +import warnings import weakref from collections import defaultdict from collections.abc import Callable @@ -27,6 +28,7 @@ from itertools import count from time import sleep from typing import Any, Literal +from unittest.mock import patch from distributed.compatibility import MACOS from distributed.scheduler import Scheduler @@ -225,7 +227,21 @@ def get_ip(): original_config = copy.deepcopy(dask.config.config) +@contextmanager +def clean_config(): + """A contextmanager that ensures the dask config is reset to the original + state during interpreter startup upon exit. + """ + with patch.dict(dask.config.config, original_config, clear=True): + yield + + def reset_config(): + warnings.warn( + "This methods will be removed soon. Consider using the contextmanager " + "distributed.utils_test.clean_config instead", + DeprecationWarning, + ) dask.config.config.clear() dask.config.config.update(copy.deepcopy(original_config)) @@ -1733,9 +1749,7 @@ def clean(threads=not WINDOWS, instances=True, timeout=1, processes=True): with check_process_leak(check=processes): with check_instances() if instances else nullcontext(): with check_active_rpc(loop, timeout): - reset_config() - dask.config.set({"distributed.comm.timeouts.connect": "5s"}) # Restore default logging levels # XXX use pytest hooks/fixtures instead? for name, level in logging_levels.items():