Fix tests that disable the shuffle extension (#7883)

dask · Jun 6, 2023 · 129b7cb · 129b7cb
1 parent 2d4ff46
commit 129b7cb
Show file tree

Hide file tree

Showing 5 changed files with 24 additions and 5 deletions.
diff --git a/distributed/http/scheduler/tests/test_semaphore_http.py b/distributed/http/scheduler/tests/test_semaphore_http.py
@@ -72,7 +72,10 @@ async def test_prometheus(c, s, a, b):
 
 
 @gen_cluster(
-    client=True, clean_kwargs={"threads": False}, scheduler_kwargs={"extensions": {}}
+    client=True,
+    clean_kwargs={"threads": False},
+    scheduler_kwargs={"extensions": {}},
+    worker_kwargs={"extensions": {}},
 )
 async def test_prometheus_without_semaphore_extension(c, s, a, b):
     pytest.importorskip("prometheus_client")

diff --git a/distributed/http/scheduler/tests/test_stealing_http.py b/distributed/http/scheduler/tests/test_stealing_http.py
@@ -95,7 +95,10 @@ async def fetch_metrics_by_cost_multipliers():
 
 
 @gen_cluster(
-    client=True, clean_kwargs={"threads": False}, scheduler_kwargs={"extensions": {}}
+    client=True,
+    clean_kwargs={"threads": False},
+    scheduler_kwargs={"extensions": {}},
+    worker_kwargs={"extensions": {}},
 )
 async def test_prometheus_without_stealing_extension(c, s, a, b):
     pytest.importorskip("prometheus_client")

diff --git a/distributed/tests/test_active_memory_manager.py b/distributed/tests/test_active_memory_manager.py
@@ -759,7 +759,11 @@ async def test_RetireWorker_amm_on_off(c, s, a, b, start_amm):
     assert "x" in b.data
 
 
-@gen_cluster(client=True, scheduler_kwargs={"extensions": {}})
+@gen_cluster(
+    client=True,
+    scheduler_kwargs={"extensions": {}},
+    worker_kwargs={"extensions": {}},
+)
 async def test_RetireWorker_no_extension(c, s, a, b):
     """retire_workers must work when the AMM extension is not loaded"""
     futures = await c.scatter({"x": 1}, workers=[a.address])

diff --git a/distributed/tests/test_spans.py b/distributed/tests/test_spans.py
@@ -103,7 +103,12 @@ async def test_multiple_tags(c, s, a, b):
     assert s.extensions["spans"].spans_search_by_tag.keys() == {"foo", "bar"}
 
 
-@gen_cluster(client=True, scheduler_kwargs={"extensions": {}})
+@gen_cluster(
+    client=True,
+    nthreads=[("", 1)],
+    scheduler_kwargs={"extensions": {}},
+    worker_kwargs={"extensions": {}},
+)
 async def test_no_extension(c, s, a, b):
     x = c.submit(inc, 1, key="x")
     assert await x == 2

diff --git a/distributed/tests/test_worker_client.py b/distributed/tests/test_worker_client.py
@@ -277,7 +277,11 @@ def test_secede_without_stealing_issue_1262():
 
     # run the loop as an inner function so all workers are closed
     # and exceptions can be examined
-    @gen_cluster(client=True, scheduler_kwargs={"extensions": {}})
+    @gen_cluster(
+        client=True,
+        scheduler_kwargs={"extensions": {}},
+        worker_kwargs={"extensions": {}},
+    )
     async def secede_test(c, s, a, b):
         def func(x):
             with worker_client() as wc: