From fbd1d1c1ef40d78c404ae51a34e82ce8329df9e4 Mon Sep 17 00:00:00 2001
From: fjetter <fjetter@users.noreply.github.com>
Date: Fri, 4 Mar 2022 15:31:31 +0100
Subject: [PATCH 1/2] Remove 5s connect timeout from clean test

---
 distributed/utils_test.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/distributed/utils_test.py b/distributed/utils_test.py
index f6625e74180..3483ec3f112 100644
--- a/distributed/utils_test.py
+++ b/distributed/utils_test.py
@@ -1735,7 +1735,6 @@ def clean(threads=not WINDOWS, instances=True, timeout=1, processes=True):
                     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():

From 3367f318f076a4a0d218408c8e901ba7287afe73 Mon Sep 17 00:00:00 2001
From: fjetter <fjetter@users.noreply.github.com>
Date: Fri, 4 Mar 2022 15:40:26 +0100
Subject: [PATCH 2/2] use contextmanager to reset dask config

---
 conftest.py               |  7 +++++++
 distributed/utils_test.py | 17 ++++++++++++++++-
 2 files changed, 23 insertions(+), 1 deletion(-)

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 3483ec3f112..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,7 +1749,6 @@ 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()
 
                         # Restore default logging levels
                         # XXX use pytest hooks/fixtures instead?