From 25514e104a4f64c1fca1d2e6f5990591925f3ae1 Mon Sep 17 00:00:00 2001
From: jmoralez <8473587+jmoralez@users.noreply.github.com>
Date: Thu, 10 Jun 2021 22:43:01 -0500
Subject: [PATCH] [dask] speed up tests (#7020)

---
 tests/python/test_with_dask.py | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/tests/python/test_with_dask.py b/tests/python/test_with_dask.py
index 1e7b7587b..a83079010 100644
--- a/tests/python/test_with_dask.py
+++ b/tests/python/test_with_dask.py
@@ -28,7 +28,6 @@ if tm.no_dask()['condition']:
     pytest.skip(msg=tm.no_dask()['reason'], allow_module_level=True)
 
 from distributed import LocalCluster, Client
-from distributed.utils_test import client, loop, cluster_fixture
 import dask.dataframe as dd
 import dask.array as da
 from xgboost.dask import DaskDMatrix
@@ -40,6 +39,20 @@ else:
     suppress = hypothesis.utils.conventions.not_set  # type:ignore
 
 
+@pytest.fixture(scope='module')
+def cluster():
+    with LocalCluster(
+        n_workers=2, threads_per_worker=2, dashboard_address=None
+    ) as dask_cluster:
+        yield dask_cluster
+
+
+@pytest.fixture
+def client(cluster):
+    with Client(cluster) as dask_client:
+        yield dask_client
+
+
 kRows = 1000
 kCols = 10
 kWorkers = 5