Use `gen_cluster` where possible in `test_dask_worker.py`
@@ -17,7 +17,7 @@ import requests
from dask.utils import tmpfile
import distributed.cli.dask_worker
-from distributed import Client, Scheduler
+from distributed import Client
from distributed.compatibility import LINUX
from distributed.deploy.utils import nprocesses_nthreads
from distributed.metrics import time
@@ -167,11 +167,8 @@ def test_no_nanny(loop):
@pytest.mark.slow
@pytest.mark.parametrize("nanny", ["--nanny", "--no-nanny"])
-@pytest.mark.asyncio
-async def test_no_reconnect(nanny):
- async with Scheduler(dashboard_address=":0") as s, Client(
- s.address, asynchronous=True
- ) as c:
+@gen_cluster(client=True, nthreads=[])
+async def test_no_reconnect(c, s, nanny):
with popen(
[
"dask-worker",
@@ -194,11 +191,8 @@ async def test_no_reconnect(nanny):
@pytest.mark.slow
@pytest.mark.parametrize("nanny", ["--nanny", "--no-nanny"])
-@pytest.mark.asyncio
-async def test_reconnect(nanny):
- async with Scheduler(dashboard_address=":0") as s, Client(
- s.address, asynchronous=True
- ) as c:
+@gen_cluster(client=True, nthreads=[])
+async def test_reconnect(c, s, nanny):
with popen(
[
"dask-worker",
@@ -532,9 +526,9 @@ async def test_integer_names(s):
assert ws.name == 123
-@pytest.mark.asyncio
@pytest.mark.parametrize("nanny", ["--nanny", "--no-nanny"])
-async def test_worker_class(cleanup, tmp_path, nanny):
+@gen_cluster(client=True, nthreads=[])
+async def test_worker_class(c, s, tmp_path, nanny):
# Create module with custom worker class
WORKER_CLASS_TEXT = """
from distributed.worker import Worker
@@ -554,8 +548,6 @@ class MyWorker(Worker):
else:
env["PYTHONPATH"] = tmpdir
- async with Scheduler(dashboard_address=":0") as s:
- async with Client(s.address, asynchronous=True) as c:
with popen(
[
"dask-worker",
