From 8bb97d26ee374d5ac3447dec5064c8dfdece7fa5 Mon Sep 17 00:00:00 2001
From: James Bourbeau <jrbourbeau@users.noreply.github.com>
Date: Wed, 8 Sep 2021 20:51:27 -0500
Subject: [PATCH] Revert "Add test setup fixture (#5242)"

This reverts commit b0c9dc4d7e5746550114a3f13831c278bc0db6d5.
---
 conftest.py               | 33 ---------------------------------
 distributed/utils_test.py | 31 +++++++++++++++++++++++++------
 2 files changed, 25 insertions(+), 39 deletions(-)

diff --git a/conftest.py b/conftest.py
index fde33e4e310..ebb9dae1457 100644
--- a/conftest.py
+++ b/conftest.py
@@ -1,11 +1,6 @@
 # https://pytest.org/latest/example/simple.html#control-skipping-of-tests-according-to-command-line-option
-import copy
-import logging
-
 import pytest
 
-import dask
-
 # Uncomment to enable more logging and checks
 # (https://docs.python.org/3/library/asyncio-dev.html)
 # Note this makes things slower and might consume much memory.
@@ -40,31 +35,3 @@ def pytest_collection_modifyitems(config, items):
 
 
 pytest_plugins = ["distributed.pytest_resourceleaks"]
-
-
-_original_config = copy.deepcopy(dask.config.config)
-# Custom preloads can interact with the test suite in unexpected ways
-# so we remove them when running tests
-for node in ["scheduler", "worker", "nanny"]:
-    _original_config["distributed"][node]["preload"] = []
-    _original_config["distributed"][node]["preload-argv"] = []
-
-_logging_levels = {
-    name: logger.level
-    for name, logger in logging.root.manager.loggerDict.items()
-    if isinstance(logger, logging.Logger)
-}
-
-
-@pytest.fixture(autouse=True)
-def initialize_test():
-
-    # Restore default logging levels
-    for name, level in _logging_levels.items():
-        logging.getLogger(name).setLevel(level)
-
-    # Ensure a clean config
-    dask.config.config.clear()
-    dask.config.config.update(copy.deepcopy(_original_config))
-
-    yield
diff --git a/distributed/utils_test.py b/distributed/utils_test.py
index b9354a5755c..a9c5496efce 100644
--- a/distributed/utils_test.py
+++ b/distributed/utils_test.py
@@ -77,6 +77,12 @@
 logger = logging.getLogger(__name__)
 
 
+logging_levels = {
+    name: logger.level
+    for name, logger in logging.root.manager.loggerDict.items()
+    if isinstance(logger, logging.Logger)
+}
+
 _TEST_TIMEOUT = 30
 _offload_executor.submit(lambda: None).result()  # create thread during import
 
@@ -209,6 +215,14 @@ def get_ip():
     remote_magic._clients.clear()
 
 
+original_config = copy.deepcopy(dask.config.config)
+
+
+def reset_config():
+    dask.config.config.clear()
+    dask.config.config.update(copy.deepcopy(original_config))
+
+
 def nodebug(func):
     """
     A decorator to disable debug facilities during timing-sensitive tests.
@@ -1571,13 +1585,18 @@ 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):
-                        with dask.config.set(
-                            {"distributed.comm.timeouts.connect": "5s"}
-                        ):
-                            yield loop
+                        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():
+                            logging.getLogger(name).setLevel(level)
+
+                        yield loop
 
-                            with suppress(AttributeError):
-                                del thread_state.on_event_loop_thread
+                        with suppress(AttributeError):
+                            del thread_state.on_event_loop_thread
 
 
 @pytest.fixture