Fix tests that disable the shuffle extension (#7883)

Author: crusaderky

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")

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")

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])

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

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: