From 0651e866fe5dc0aadd3fe3af54d589f8763f86fd Mon Sep 17 00:00:00 2001
From: Thomas Grainger <tagrain@gmail.com>
Date: Fri, 12 Aug 2022 14:03:30 +0100
Subject: [PATCH 1/3] use async with Client in tests

---
 distributed/deploy/tests/test_adaptive.py |  93 ++--
 distributed/deploy/tests/test_local.py    |  64 ++-
 distributed/tests/test_client.py          | 567 ++++++++++------------
 distributed/tests/test_failed_workers.py  |  56 ++-
 distributed/tests/test_publish.py         | 100 ++--
 distributed/tests/test_queues.py          |  62 ++-
 distributed/tests/test_scheduler.py       |  13 +-
 distributed/tests/test_utils_test.py      |   7 +-
 distributed/tests/test_variable.py        |  62 ++-
 distributed/tests/test_worker.py          |   5 +-
 distributed/utils_test.py                 | 159 +++---
 11 files changed, 572 insertions(+), 616 deletions(-)

diff --git a/distributed/deploy/tests/test_adaptive.py b/distributed/deploy/tests/test_adaptive.py
index 1248db0f0d8..04b518b063b 100644
--- a/distributed/deploy/tests/test_adaptive.py
+++ b/distributed/deploy/tests/test_adaptive.py
@@ -117,50 +117,47 @@ def scale_up(self, n, **kwargs):
 
 @gen_test()
 async def test_min_max():
-    cluster = await LocalCluster(
+    async with LocalCluster(
         n_workers=0,
         silence_logs=False,
         processes=False,
         dashboard_address=":0",
         asynchronous=True,
         threads_per_worker=1,
-    )
-    try:
+    ) as cluster:
         adapt = cluster.adapt(minimum=1, maximum=2, interval="20 ms", wait_count=10)
-        c = await Client(cluster, asynchronous=True)
-
-        start = time()
-        while not cluster.scheduler.workers:
-            await asyncio.sleep(0.01)
-            assert time() < start + 1
-
-        await asyncio.sleep(0.2)
-        assert len(cluster.scheduler.workers) == 1
-        assert len(adapt.log) == 1 and adapt.log[-1][1] == {"status": "up", "n": 1}
+        async with Client(cluster, asynchronous=True) as c:
+            start = time()
+            while not cluster.scheduler.workers:
+                await asyncio.sleep(0.01)
+                assert time() < start + 1
 
-        futures = c.map(slowinc, range(100), delay=0.1)
+            await asyncio.sleep(0.2)
+            assert len(cluster.scheduler.workers) == 1
+            assert len(adapt.log) == 1 and adapt.log[-1][1] == {"status": "up", "n": 1}
 
-        start = time()
-        while len(cluster.scheduler.workers) < 2:
-            await asyncio.sleep(0.01)
-            assert time() < start + 1
+            futures = c.map(slowinc, range(100), delay=0.1)
 
-        assert len(cluster.scheduler.workers) == 2
-        await asyncio.sleep(0.5)
-        assert len(cluster.scheduler.workers) == 2
-        assert len(cluster.workers) == 2
-        assert len(adapt.log) == 2 and all(d["status"] == "up" for _, d in adapt.log)
+            start = time()
+            while len(cluster.scheduler.workers) < 2:
+                await asyncio.sleep(0.01)
+                assert time() < start + 1
+
+            assert len(cluster.scheduler.workers) == 2
+            await asyncio.sleep(0.5)
+            assert len(cluster.scheduler.workers) == 2
+            assert len(cluster.workers) == 2
+            assert len(adapt.log) == 2 and all(
+                d["status"] == "up" for _, d in adapt.log
+            )
 
-        del futures
+            del futures
 
-        start = time()
-        while len(cluster.scheduler.workers) != 1:
-            await asyncio.sleep(0.01)
-            assert time() < start + 2
-        assert adapt.log[-1][1]["status"] == "down"
-    finally:
-        await c.close()
-        await cluster.close()
+            start = time()
+            while len(cluster.scheduler.workers) != 1:
+                await asyncio.sleep(0.01)
+                assert time() < start + 2
+            assert adapt.log[-1][1]["status"] == "down"
 
 
 @gen_test()
@@ -194,16 +191,14 @@ async def test_adapt_quickly():
     Instead we want to wait a few beats before removing a worker in case the
     user is taking a brief pause between work
     """
-    cluster = await LocalCluster(
+    async with LocalCluster(
         n_workers=0,
         asynchronous=True,
         processes=False,
         silence_logs=False,
         dashboard_address=":0",
-    )
-    client = await Client(cluster, asynchronous=True)
-    adapt = cluster.adapt(interval="20 ms", wait_count=5, maximum=10)
-    try:
+    ) as cluster, Client(cluster, asynchronous=True) as client:
+        adapt = cluster.adapt(interval="20 ms", wait_count=5, maximum=10)
         future = client.submit(slowinc, 1, delay=0.100)
         await wait(future)
         assert len(adapt.log) == 1
@@ -241,9 +236,6 @@ async def test_adapt_quickly():
 
         await asyncio.sleep(0.1)
         assert len(cluster.workers) == 1
-    finally:
-        await client.close()
-        await cluster.close()
 
 
 @gen_test()
@@ -255,20 +247,19 @@ async def test_adapt_down():
         processes=False,
         silence_logs=False,
         dashboard_address=":0",
-    ) as cluster:
-        async with Client(cluster, asynchronous=True) as client:
-            cluster.adapt(interval="20ms", maximum=5)
+    ) as cluster, Client(cluster, asynchronous=True) as client:
+        cluster.adapt(interval="20ms", maximum=5)
 
-            futures = client.map(slowinc, range(1000), delay=0.1)
-            while len(cluster.scheduler.workers) < 5:
-                await asyncio.sleep(0.1)
+        futures = client.map(slowinc, range(1000), delay=0.1)
+        while len(cluster.scheduler.workers) < 5:
+            await asyncio.sleep(0.1)
 
-            cluster.adapt(maximum=2)
+        cluster.adapt(maximum=2)
 
-            start = time()
-            while len(cluster.scheduler.workers) != 2:
-                await asyncio.sleep(0.1)
-                assert time() < start + 60
+        start = time()
+        while len(cluster.scheduler.workers) != 2:
+            await asyncio.sleep(0.1)
+            assert time() < start + 60
 
 
 @gen_test()
diff --git a/distributed/deploy/tests/test_local.py b/distributed/deploy/tests/test_local.py
index 1da51d9f4e4..bc64cfca0b6 100644
--- a/distributed/deploy/tests/test_local.py
+++ b/distributed/deploy/tests/test_local.py
@@ -5,11 +5,11 @@
 import sys
 from threading import Lock
 from time import sleep
+from unittest import mock
 from urllib.parse import urlparse
 
 import pytest
 from tornado.httpclient import AsyncHTTPClient
-from tornado.ioloop import IOLoop
 
 from dask.system import CPU_COUNT
 
@@ -245,19 +245,21 @@ def test_Client_solo(loop):
 @gen_test()
 async def test_duplicate_clients():
     pytest.importorskip("bokeh")
-    c1 = await Client(
+    async with Client(
         processes=False, silence_logs=False, dashboard_address=9876, asynchronous=True
-    )
-    with pytest.warns(Warning) as info:
-        c2 = await Client(
-            processes=False,
-            silence_logs=False,
-            dashboard_address=9876,
-            asynchronous=True,
-        )
-
-    assert "dashboard" in c1.cluster.scheduler.services
-    assert "dashboard" in c2.cluster.scheduler.services
+    ) as c1:
+        c1_services = c1.cluster.scheduler.services
+        with pytest.warns(Warning) as info:
+            async with Client(
+                processes=False,
+                silence_logs=False,
+                dashboard_address=9876,
+                asynchronous=True,
+            ) as c2:
+                c2_services = c2.cluster.scheduler.services
+
+    assert c1_services == {"dashboard": mock.ANY}
+    assert c2_services == {"dashboard": mock.ANY}
 
     assert any(
         all(
@@ -266,8 +268,6 @@ async def test_duplicate_clients():
         )
         for msg in info.list
     )
-    await c1.close()
-    await c2.close()
 
 
 def test_Client_kwargs(loop):
@@ -824,35 +824,29 @@ class MyCluster(LocalCluster):
         def scale_down(self, *args, **kwargs):
             pass
 
-    loop = IOLoop.current()
-    cluster = await MyCluster(
+    async with MyCluster(
         n_workers=0,
         processes=False,
         silence_logs=False,
         dashboard_address=":0",
-        loop=loop,
+        loop=None,
         asynchronous=True,
-    )
-    c = await Client(cluster, asynchronous=True)
-
-    assert not cluster.workers
-
-    await cluster.scale(2)
+    ) as cluster, Client(cluster, asynchronous=True) as c:
+        assert not cluster.workers
 
-    start = time()
-    while len(cluster.scheduler.workers) != 2:
-        await asyncio.sleep(0.01)
-        assert time() < start + 3
+        await cluster.scale(2)
 
-    await cluster.scale(1)
+        start = time()
+        while len(cluster.scheduler.workers) != 2:
+            await asyncio.sleep(0.01)
+            assert time() < start + 3
 
-    start = time()
-    while len(cluster.scheduler.workers) != 1:
-        await asyncio.sleep(0.01)
-        assert time() < start + 3
+        await cluster.scale(1)
 
-    await c.close()
-    await cluster.close()
+        start = time()
+        while len(cluster.scheduler.workers) != 1:
+            await asyncio.sleep(0.01)
+            assert time() < start + 3
 
 
 def test_local_tls_restart(loop):
diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
index fbf2af7f99d..334b43f3fdd 100644
--- a/distributed/tests/test_client.py
+++ b/distributed/tests/test_client.py
@@ -6,6 +6,7 @@
 import gc
 import inspect
 import logging
+import operator
 import os
 import pathlib
 import pickle
@@ -534,17 +535,14 @@ async def test_exceptions(c, s, a, b):
 
 @gen_cluster()
 async def test_gc(s, a, b):
-    c = await Client(s.address, asynchronous=True)
-
-    x = c.submit(inc, 10)
-    await x
-    assert s.tasks[x.key].who_has
-    x.__del__()
-    await async_wait_for(
-        lambda: x.key not in s.tasks or not s.tasks[x.key].who_has, timeout=0.3
-    )
-
-    await c.close()
+    async with Client(s.address, asynchronous=True) as c:
+        x = c.submit(inc, 10)
+        await x
+        assert s.tasks[x.key].who_has
+        x.__del__()
+        await async_wait_for(
+            lambda: x.key not in s.tasks or not s.tasks[x.key].who_has, timeout=0.3
+        )
 
 
 def test_thread(c):
@@ -587,13 +585,12 @@ async def test_gather(c, s, a, b):
 
 @gen_cluster(client=True)
 async def test_gather_mismatched_client(c, s, a, b):
-    c2 = await Client(s.address, asynchronous=True)
+    async with Client(s.address, asynchronous=True) as c2:
+        x = c.submit(inc, 10)
+        y = c2.submit(inc, 5)
 
-    x = c.submit(inc, 10)
-    y = c2.submit(inc, 5)
-
-    with pytest.raises(ValueError, match="Futures created by another client"):
-        await c.gather([x, y])
+        with pytest.raises(ValueError, match="Futures created by another client"):
+            await c.gather([x, y])
 
 
 @gen_cluster(client=True)
@@ -1067,20 +1064,16 @@ async def test_map_quotes(c, s, a, b):
 
 @gen_cluster()
 async def test_two_consecutive_clients_share_results(s, a, b):
-    c = await Client(s.address, asynchronous=True)
-
-    x = c.submit(random.randint, 0, 1000, pure=True)
-    xx = await x
-
-    f = await Client(s.address, asynchronous=True)
+    async with Client(s.address, asynchronous=True) as c:
 
-    y = f.submit(random.randint, 0, 1000, pure=True)
-    yy = await y
+        x = c.submit(random.randint, 0, 1000, pure=True)
+        xx = await x
 
-    assert xx == yy
+        async with Client(s.address, asynchronous=True) as f:
+            y = f.submit(random.randint, 0, 1000, pure=True)
+            yy = await y
 
-    await c.close()
-    await f.close()
+            assert xx == yy
 
 
 @gen_cluster(client=True)
@@ -1423,14 +1416,10 @@ async def test_scatter_direct(c, s, a, b):
 
 @gen_cluster()
 async def test_scatter_direct_2(s, a, b):
-    c = await Client(s.address, asynchronous=True, heartbeat_interval=10)
-
-    last = s.clients[c.id].last_seen
-
-    while s.clients[c.id].last_seen == last:
-        await asyncio.sleep(0.10)
-
-    await c.close()
+    async with Client(s.address, asynchronous=True, heartbeat_interval=10) as c:
+        last = s.clients[c.id].last_seen
+        while s.clients[c.id].last_seen == last:
+            await asyncio.sleep(0.10)
 
 
 @gen_cluster(client=True)
@@ -2014,21 +2003,26 @@ async def test_badly_serialized_input_stderr(capsys, c):
     assert future.status == "error"
 
 
-def test_repr(loop):
-    funcs = [str, repr, lambda x: x._repr_html_()]
+@pytest.mark.parametrize(
+    "func",
+    [
+        str,
+        repr,
+        operator.methodcaller("_repr_html_"),
+    ],
+)
+def test_repr(loop, func):
     with cluster(nworkers=3, worker_kwargs={"memory_limit": "2 GiB"}) as (s, [a, b, c]):
         with Client(s["address"], loop=loop) as c:
-            for func in funcs:
-                text = func(c)
-                assert c.scheduler.address in text
-                assert "threads=3" in text or "Total threads: </strong>" in text
-                assert "6.00 GiB" in text
-                if "<table" not in text:
-                    assert len(text) < 80
-
-        for func in funcs:
             text = func(c)
-            assert "No scheduler connected" in text
+            assert c.scheduler.address in text
+            assert "threads=3" in text or "Total threads: </strong>" in text
+            assert "6.00 GiB" in text
+            if "<table" not in text:
+                assert len(text) < 80
+
+        text = func(c)
+        assert "No scheduler connected" in text
 
 
 @gen_cluster(client=True)
@@ -2043,17 +2037,12 @@ async def test_repr_no_memory_limit(c, s, a, b):
 
 @gen_test()
 async def test_repr_localcluster():
-    cluster = await LocalCluster(
+    async with LocalCluster(
         processes=False, dashboard_address=":0", asynchronous=True
-    )
-    client = await Client(cluster, asynchronous=True)
-    try:
+    ) as cluster, Client(cluster, asynchronous=True) as client:
         text = client._repr_html_()
         assert cluster.scheduler.address in text
         assert is_valid_xml(client._repr_html_())
-    finally:
-        await client.close()
-        await cluster.close()
 
 
 @gen_cluster(client=True)
@@ -2176,32 +2165,28 @@ def test_repr_sync(c):
 
 @gen_cluster()
 async def test_multi_client(s, a, b):
-    c = await Client(s.address, asynchronous=True)
-    f = await Client(s.address, asynchronous=True)
-
-    assert set(s.client_comms) == {c.id, f.id}
+    async with Client(s.address, asynchronous=True) as f:
+        async with Client(s.address, asynchronous=True) as c:
 
-    x = c.submit(inc, 1)
-    y = f.submit(inc, 2)
-    y2 = c.submit(inc, 2)
+            assert set(s.client_comms) == {c.id, f.id}
 
-    assert y.key == y2.key
+            x = c.submit(inc, 1)
+            y = f.submit(inc, 2)
+            y2 = c.submit(inc, 2)
 
-    await wait([x, y])
+            assert y.key == y2.key
 
-    assert {ts.key for ts in s.clients[c.id].wants_what} == {x.key, y.key}
-    assert {ts.key for ts in s.clients[f.id].wants_what} == {y.key}
+            await wait([x, y])
 
-    await c.close()
+            assert {ts.key for ts in s.clients[c.id].wants_what} == {x.key, y.key}
+            assert {ts.key for ts in s.clients[f.id].wants_what} == {y.key}
 
-    while c.id in s.clients:
-        await asyncio.sleep(0.01)
-
-    assert c.id not in s.clients
-    assert c.id not in s.tasks[y.key].who_wants
-    assert x.key not in s.tasks
+        while c.id in s.clients:
+            await asyncio.sleep(0.01)
 
-    await f.close()
+        assert c.id not in s.clients
+        assert c.id not in s.tasks[y.key].who_wants
+        assert x.key not in s.tasks
 
     while s.tasks:
         await asyncio.sleep(0.01)
@@ -2234,42 +2219,39 @@ async def test_cleanup_after_broken_client_connection(s, a, b):
 
 @gen_cluster()
 async def test_multi_garbage_collection(s, a, b):
-    c = await Client(s.address, asynchronous=True)
-
-    f = await Client(s.address, asynchronous=True)
+    async with Client(s.address, asynchronous=True) as c, Client(
+        s.address, asynchronous=True
+    ) as f:
 
-    x = c.submit(inc, 1)
-    y = f.submit(inc, 2)
-    y2 = c.submit(inc, 2)
-
-    assert y.key == y2.key
+        x = c.submit(inc, 1)
+        y = f.submit(inc, 2)
+        y2 = c.submit(inc, 2)
 
-    await wait([x, y])
+        assert y.key == y2.key
 
-    x.__del__()
-    while x.key in a.data or x.key in b.data:
-        await asyncio.sleep(0.01)
+        await wait([x, y])
 
-    assert {ts.key for ts in s.clients[c.id].wants_what} == {y.key}
-    assert {ts.key for ts in s.clients[f.id].wants_what} == {y.key}
+        x.__del__()
+        while x.key in a.data or x.key in b.data:
+            await asyncio.sleep(0.01)
 
-    y.__del__()
-    while x.key in {ts.key for ts in s.clients[f.id].wants_what}:
-        await asyncio.sleep(0.01)
+        assert {ts.key for ts in s.clients[c.id].wants_what} == {y.key}
+        assert {ts.key for ts in s.clients[f.id].wants_what} == {y.key}
 
-    await asyncio.sleep(0.1)
-    assert y.key in a.data or y.key in b.data
-    assert {ts.key for ts in s.clients[c.id].wants_what} == {y.key}
-    assert not s.clients[f.id].wants_what
+        y.__del__()
+        while x.key in {ts.key for ts in s.clients[f.id].wants_what}:
+            await asyncio.sleep(0.01)
 
-    y2.__del__()
-    while y.key in a.data or y.key in b.data:
-        await asyncio.sleep(0.01)
+        await asyncio.sleep(0.1)
+        assert y.key in a.data or y.key in b.data
+        assert {ts.key for ts in s.clients[c.id].wants_what} == {y.key}
+        assert not s.clients[f.id].wants_what
 
-    assert not s.tasks
+        y2.__del__()
+        while y.key in a.data or y.key in b.data:
+            await asyncio.sleep(0.01)
 
-    await c.close()
-    await f.close()
+        assert not s.tasks
 
 
 @gen_cluster(client=True)
@@ -2868,11 +2850,11 @@ def test_client_num_fds(loop):
 
 @gen_cluster()
 async def test_startup_close_startup(s, a, b):
-    c = await Client(s.address, asynchronous=True)
-    await c.close()
+    async with Client(s.address, asynchronous=True):
+        pass
 
-    c = await Client(s.address, asynchronous=True)
-    await c.close()
+    async with Client(s.address, asynchronous=True):
+        pass
 
 
 @pytest.mark.filterwarnings("ignore:There is no current event loop:DeprecationWarning")
@@ -3570,11 +3552,10 @@ def test_as_completed_next_batch(c):
 
 @gen_cluster(nthreads=[])
 async def test_status(s):
-    c = await Client(s.address, asynchronous=True)
-    assert c.status == "running"
-    x = c.submit(inc, 1)
+    async with Client(s.address, asynchronous=True) as c:
+        assert c.status == "running"
+        x = c.submit(inc, 1)
 
-    await c.close()
     assert c.status == "closed"
 
 
@@ -3627,54 +3608,52 @@ async def hard_stop(s):
         s.stop()
         await Server.close(s)
 
-    s = await Scheduler(port=port)
-    w = await Worker(f"127.0.0.1:{port}")
-    c = await Client(f"127.0.0.1:{port}", asynchronous=True)
-    await c.wait_for_workers(1, timeout=10)
-    x = c.submit(inc, 1)
-    assert (await x) == 2
-    await hard_stop(s)
-
-    start = time()
-    while c.status != "connecting":
-        assert time() < start + 10
-        await asyncio.sleep(0.01)
-
-    assert x.status == "cancelled"
-    with pytest.raises(CancelledError):
-        await x
-
-    s = await Scheduler(port=port)
-    start = time()
-    while c.status != "running":
-        await asyncio.sleep(0.1)
-        assert time() < start + 10
-
-    await w.finished()
-    w = await Worker(f"127.0.0.1:{port}")
-
-    start = time()
-    while len(await c.nthreads()) != 1:
-        await asyncio.sleep(0.05)
-        assert time() < start + 10
-
-    x = c.submit(inc, 1)
-    assert (await x) == 2
-    await hard_stop(s)
-
-    start = time()
-    while True:
-        assert time() < start + 10
-        try:
-            await x
-            assert False
-        except CommClosedError:
-            continue
-        except CancelledError:
-            break
-
-    await w.close()
-    await c._close(fast=True)
+    async with Scheduler(port=port) as s:
+        async with Client(f"127.0.0.1:{port}", asynchronous=True) as c:
+            async with Worker(f"127.0.0.1:{port}") as w:
+                await c.wait_for_workers(1, timeout=10)
+                x = c.submit(inc, 1)
+                assert (await x) == 2
+                await hard_stop(s)
+
+                start = time()
+                while c.status != "connecting":
+                    assert time() < start + 10
+                    await asyncio.sleep(0.01)
+
+                assert x.status == "cancelled"
+                with pytest.raises(CancelledError):
+                    await x
+
+                async with Scheduler(port=port) as s2:
+                    start = time()
+                    while c.status != "running":
+                        await asyncio.sleep(0.1)
+                        assert time() < start + 10
+
+                    await w.finished()
+                    w = await Worker(f"127.0.0.1:{port}")
+
+                    start = time()
+                    while len(await c.nthreads()) != 1:
+                        await asyncio.sleep(0.05)
+                        assert time() < start + 10
+
+                    x = c.submit(inc, 1)
+                    assert (await x) == 2
+                    await hard_stop(s2)
+
+                start = time()
+                while True:
+                    assert time() < start + 10
+                    try:
+                        await x
+                        assert False
+                    except CommClosedError:
+                        continue
+                    except CancelledError:
+                        break
+            await c._close(fast=True)
 
 
 class UnhandledExceptions(Exception):
@@ -3796,46 +3775,44 @@ async def start_worker(sleep, duration, repeat=1):
 
 @gen_cluster()
 async def test_idempotence(s, a, b):
-    c = await Client(s.address, asynchronous=True)
-    f = await Client(s.address, asynchronous=True)
+    async with Client(s.address, asynchronous=True) as c, Client(
+        s.address, asynchronous=True
+    ) as f:
 
-    # Submit
-    x = c.submit(inc, 1)
-    await x
-    log = list(s.transition_log)
-
-    len_single_submit = len(log)  # see last assert
-
-    y = f.submit(inc, 1)
-    assert x.key == y.key
-    await y
-    await asyncio.sleep(0.1)
-    log2 = list(s.transition_log)
-    assert log == log2
-
-    # Error
-    a = c.submit(div, 1, 0)
-    await wait(a)
-    assert a.status == "error"
-    log = list(s.transition_log)
+        # Submit
+        x = c.submit(inc, 1)
+        await x
+        log = list(s.transition_log)
 
-    b = f.submit(div, 1, 0)
-    assert a.key == b.key
-    await wait(b)
-    await asyncio.sleep(0.1)
-    log2 = list(s.transition_log)
-    assert log == log2
+        len_single_submit = len(log)  # see last assert
 
-    s.transition_log.clear()
-    # Simultaneous Submit
-    d = c.submit(inc, 2)
-    e = c.submit(inc, 2)
-    await wait([d, e])
+        y = f.submit(inc, 1)
+        assert x.key == y.key
+        await y
+        await asyncio.sleep(0.1)
+        log2 = list(s.transition_log)
+        assert log == log2
+
+        # Error
+        a = c.submit(div, 1, 0)
+        await wait(a)
+        assert a.status == "error"
+        log = list(s.transition_log)
+
+        b = f.submit(div, 1, 0)
+        assert a.key == b.key
+        await wait(b)
+        await asyncio.sleep(0.1)
+        log2 = list(s.transition_log)
+        assert log == log2
 
-    assert len(s.transition_log) == len_single_submit
+        s.transition_log.clear()
+        # Simultaneous Submit
+        d = c.submit(inc, 2)
+        e = c.submit(inc, 2)
+        await wait([d, e])
 
-    await c.close()
-    await f.close()
+        assert len(s.transition_log) == len_single_submit
 
 
 def test_scheduler_info(c):
@@ -4010,60 +3987,56 @@ async def test_scatter_compute_store_lose_processing(c, s, a, b):
 
 @gen_cluster()
 async def test_serialize_future(s, a, b):
-    c1 = await Client(s.address, asynchronous=True)
-    c2 = await Client(s.address, asynchronous=True)
-
-    future = c1.submit(lambda: 1)
-    result = await future
+    async with Client(s.address, asynchronous=True) as c1, Client(
+        s.address, asynchronous=True
+    ) as c2:
 
-    for ci in (c1, c2):
-        for ctxman in lambda ci: ci.as_current(), lambda ci: temp_default_client(ci):
-            with ctxman(ci):
-                future2 = pickle.loads(pickle.dumps(future))
-                assert future2.client is ci
-                assert stringify(future2.key) in ci.futures
-                result2 = await future2
-                assert result == result2
+        future = c1.submit(lambda: 1)
+        result = await future
 
-    await c1.close()
-    await c2.close()
+        for ci in (c1, c2):
+            for ctxman in lambda ci: ci.as_current(), lambda ci: temp_default_client(
+                ci
+            ):
+                with ctxman(ci):
+                    future2 = pickle.loads(pickle.dumps(future))
+                    assert future2.client is ci
+                    assert stringify(future2.key) in ci.futures
+                    result2 = await future2
+                    assert result == result2
 
 
 @gen_cluster()
 async def test_temp_default_client(s, a, b):
-    c1 = await Client(s.address, asynchronous=True)
-    c2 = await Client(s.address, asynchronous=True)
-
-    with temp_default_client(c1):
-        assert default_client() is c1
-        assert default_client(c2) is c2
+    async with Client(s.address, asynchronous=True) as c1, Client(
+        s.address, asynchronous=True
+    ) as c2:
 
-    with temp_default_client(c2):
-        assert default_client() is c2
-        assert default_client(c1) is c1
+        with temp_default_client(c1):
+            assert default_client() is c1
+            assert default_client(c2) is c2
 
-    await c1.close()
-    await c2.close()
+        with temp_default_client(c2):
+            assert default_client() is c2
+            assert default_client(c1) is c1
 
 
 @gen_cluster(client=True)
 async def test_as_current(c, s, a, b):
-    c1 = await Client(s.address, asynchronous=True)
-    c2 = await Client(s.address, asynchronous=True)
-
-    with temp_default_client(c):
-        assert Client.current() is c
-        with pytest.raises(ValueError):
-            Client.current(allow_global=False)
-        with c1.as_current():
-            assert Client.current() is c1
-            assert Client.current(allow_global=True) is c1
-        with c2.as_current():
-            assert Client.current() is c2
-            assert Client.current(allow_global=True) is c2
-
-    await c1.close()
-    await c2.close()
+    async with Client(s.address, asynchronous=True) as c1, Client(
+        s.address, asynchronous=True
+    ) as c2:
+
+        with temp_default_client(c):
+            assert Client.current() is c
+            with pytest.raises(ValueError):
+                Client.current(allow_global=False)
+            with c1.as_current():
+                assert Client.current() is c1
+                assert Client.current(allow_global=True) is c1
+            with c2.as_current():
+                assert Client.current() is c2
+                assert Client.current(allow_global=True) is c2
 
 
 def test_as_current_is_thread_local(s, loop):
@@ -4563,14 +4536,27 @@ async def test_client_timeout():
     straight away
     """
     port = open_port()
-    with dask.config.set({"distributed.comm.timeouts.connect": "10s"}):
-        c = Client(f"127.0.0.1:{port}", asynchronous=True)
-        client_start_fut = asyncio.ensure_future(c)
+    stop_event = asyncio.Event()
+
+    async def run_client():
+        try:
+            async with Client(f"127.0.0.1:{port}", asynchronous=True) as c:
+                return await c.run_on_scheduler(lambda: 123)
+        finally:
+            stop_event.set()
+
+    async def run_scheduler_after_2_seconds():
+        # TODO: start a scheduler that waits for the first connection and
+        # closes it
         await asyncio.sleep(2)
         async with Scheduler(port=port, dashboard_address=":0"):
-            await client_start_fut
-            assert await c.run_on_scheduler(lambda: 123) == 123
-            await c.close()
+            await stop_event.wait()
+
+    with dask.config.set({"distributed.comm.timeouts.connect": "10s"}):
+        assert await asyncio.gather(
+            run_client(),
+            run_scheduler_after_2_seconds(),
+        ) == [123, None]
 
 
 @gen_cluster(client=True)
@@ -4955,14 +4941,14 @@ def test_quiet_client_close_when_cluster_is_closed_before_client(loop):
 
 @gen_cluster()
 async def test_close(s, a, b):
-    c = await Client(s.address, asynchronous=True)
-    future = c.submit(inc, 1)
-    await wait(future)
-    assert c.id in s.clients
-    await c.close()
+    async with Client(s.address, asynchronous=True) as c:
+        future = c.submit(inc, 1)
+        await wait(future)
+        assert c.id in s.clients
+        await c.close()
 
-    while c.id in s.clients or s.tasks:
-        await asyncio.sleep(0.01)
+        while c.id in s.clients or s.tasks:
+            await asyncio.sleep(0.01)
 
 
 def test_threadsafe(c):
@@ -5499,9 +5485,8 @@ async def test_profile_keys(c, s, a, b):
 @gen_cluster()
 async def test_client_with_name(s, a, b):
     with captured_logger("distributed.scheduler") as sio:
-        client = await Client(s.address, asynchronous=True, name="foo")
-        assert "foo" in client.id
-        await client.close()
+        async with Client(s.address, asynchronous=True, name="foo") as client:
+            assert "foo" in client.id
 
     text = sio.getvalue()
     assert "foo" in text
@@ -5521,13 +5506,11 @@ async def test_future_auto_inform(c, s, a, b):
     x = c.submit(inc, 1)
     await wait(x)
 
-    client = await Client(s.address, asynchronous=True)
-    future = Future(x.key, client)
-
-    while future.status != "finished":
-        await asyncio.sleep(0.01)
+    async with Client(s.address, asynchronous=True) as client:
+        future = Future(x.key, client)
 
-    await client.close()
+        while future.status != "finished":
+            await asyncio.sleep(0.01)
 
 
 @pytest.mark.filterwarnings("ignore:There is no current event loop:DeprecationWarning")
@@ -5644,13 +5627,10 @@ async def test_avoid_delayed_finalize(c, s, a, b):
 async def test_config_scheduler_address(s, a, b):
     with dask.config.set({"scheduler-address": s.address}):
         with captured_logger("distributed.client") as sio:
-            c = await Client(asynchronous=True)
-            assert c.scheduler.address == s.address
+            async with Client(asynchronous=True) as c:
+                assert c.scheduler.address == s.address
 
-        text = sio.getvalue()
-        assert s.address in text
-
-        await c.close()
+    assert sio.getvalue() == f"Config value `scheduler-address` found: {s.address}\n"
 
 
 @gen_cluster(client=True)
@@ -5681,10 +5661,8 @@ async def test_unhashable_function(c, s, a, b):
 @gen_cluster()
 async def test_client_name(s, a, b):
     with dask.config.set({"client-name": "hello-world"}):
-        c = await Client(s.address, asynchronous=True)
-        assert any("hello-world" in name for name in list(s.clients))
-
-    await c.close()
+        async with Client(s.address, asynchronous=True) as c:
+            assert any("hello-world" in name for name in list(s.clients))
 
 
 def test_client_doesnt_close_given_loop(loop_in_thread, s, a, b):
@@ -5731,12 +5709,10 @@ async def test_client_timeout_2():
         start = time()
         c = Client(f"127.0.0.1:{port}", asynchronous=True)
         with pytest.raises((TimeoutError, IOError)):
-            await c
+            async with c:
+                pass
         stop = time()
-
         assert c.status == "closed"
-        await c.close()
-
         assert stop - start < 1
 
 
@@ -5751,8 +5727,8 @@ async def test_client_active_bad_port():
     with dask.config.set({"distributed.comm.timeouts.connect": "10ms"}):
         c = Client("127.0.0.1:8080", asynchronous=True)
         with pytest.raises((TimeoutError, IOError)):
-            await c
-        await c._close(fast=True)
+            async with c:
+                pass
     http_server.stop()
 
 
@@ -5794,50 +5770,45 @@ async def test_turn_off_pickle(c, s, a, b, direct):
 async def test_de_serialization(s, a, b):
     np = pytest.importorskip("numpy")
 
-    c = await Client(
+    async with Client(
         s.address,
         asynchronous=True,
         serializers=["msgpack", "pickle"],
         deserializers=["msgpack"],
-    )
-    try:
+    ) as c:
         # Can send complex data
         future = await c.scatter(np.ones(5))
 
         # But can not retrieve it
         with pytest.raises(TypeError):
             result = await future
-    finally:
-        await c.close()
 
 
 @gen_cluster()
 async def test_de_serialization_none(s, a, b):
     np = pytest.importorskip("numpy")
 
-    c = await Client(s.address, asynchronous=True, deserializers=["msgpack"])
-    try:
+    async with Client(s.address, asynchronous=True, deserializers=["msgpack"]) as c:
         # Can send complex data
         future = await c.scatter(np.ones(5))
 
         # But can not retrieve it
         with pytest.raises(TypeError):
             result = await future
-    finally:
-        await c.close()
 
 
 @gen_cluster()
 async def test_client_repr_closed(s, a, b):
-    c = await Client(s.address, asynchronous=True)
-    await c.close()
-    c._repr_html_()
+    async with Client(s.address, asynchronous=True) as c:
+        pass
+    assert "No scheduler connected." in c._repr_html_()
 
 
+@pytest.mark.skip
 def test_client_repr_closed_sync(loop):
     with Client(loop=loop, processes=False, dashboard_address=":0") as c:
         pass
-    c._repr_html_()
+    assert "No scheduler connected." in c._repr_html_()
 
 
 @pytest.mark.xfail(reason="https://github.com/dask/dask/pull/6807")
@@ -5917,13 +5888,11 @@ def test_no_threads_lingering():
 
 @gen_cluster()
 async def test_direct_async(s, a, b):
-    c = await Client(s.address, asynchronous=True, direct_to_workers=True)
-    assert c.direct_to_workers
-    await c.close()
+    async with Client(s.address, asynchronous=True, direct_to_workers=True) as c:
+        assert c.direct_to_workers
 
-    c = await Client(s.address, asynchronous=True, direct_to_workers=False)
-    assert not c.direct_to_workers
-    await c.close()
+    async with Client(s.address, asynchronous=True, direct_to_workers=False) as c:
+        assert not c.direct_to_workers
 
 
 def test_direct_sync(c):
@@ -5937,17 +5906,15 @@ def f():
 
 @gen_cluster()
 async def test_mixing_clients(s, a, b):
-    c1 = await Client(s.address, asynchronous=True)
-    c2 = await Client(s.address, asynchronous=True)
-
-    future = c1.submit(inc, 1)
-    with pytest.raises(ValueError):
-        c2.submit(inc, future)
+    async with Client(s.address, asynchronous=True) as c1, Client(
+        s.address, asynchronous=True
+    ) as c2:
 
-    assert not c2.futures  # Don't create Futures on second Client
+        future = c1.submit(inc, 1)
+        with pytest.raises(ValueError):
+            c2.submit(inc, future)
 
-    await c1.close()
-    await c2.close()
+        assert not c2.futures  # Don't create Futures on second Client
 
 
 @gen_cluster(client=True)
diff --git a/distributed/tests/test_failed_workers.py b/distributed/tests/test_failed_workers.py
index 7ebdbc4ebd3..48bb30af592 100644
--- a/distributed/tests/test_failed_workers.py
+++ b/distributed/tests/test_failed_workers.py
@@ -175,43 +175,41 @@ def test_worker_doesnt_await_task_completion(loop):
 
 @gen_cluster(Worker=Nanny, timeout=60)
 async def test_multiple_clients_restart(s, a, b):
-    c1 = await Client(s.address, asynchronous=True)
-    c2 = await Client(s.address, asynchronous=True)
+    async with Client(s.address, asynchronous=True) as c1, Client(
+        s.address, asynchronous=True
+    ) as c2:
 
-    x = c1.submit(inc, 1)
-    y = c2.submit(inc, 2)
-    xx = await x
-    yy = await y
-    assert xx == 2
-    assert yy == 3
+        x = c1.submit(inc, 1)
+        y = c2.submit(inc, 2)
+        xx = await x
+        yy = await y
+        assert xx == 2
+        assert yy == 3
 
-    await c1.restart()
+        await c1.restart()
 
-    assert x.cancelled()
-    start = time()
-    while not y.cancelled():
-        await asyncio.sleep(0.01)
-        assert time() < start + 5
-
-    assert not c1.futures
-    assert not c2.futures
+        assert x.cancelled()
+        start = time()
+        while not y.cancelled():
+            await asyncio.sleep(0.01)
+            assert time() < start + 5
 
-    # Ensure both clients still work after restart.
-    # Reusing a previous key has no effect.
-    x2 = c1.submit(inc, 1, key=x.key)
-    y2 = c2.submit(inc, 2, key=y.key)
+        assert not c1.futures
+        assert not c2.futures
 
-    assert x2._generation != x._generation
-    assert y2._generation != y._generation
+        # Ensure both clients still work after restart.
+        # Reusing a previous key has no effect.
+        x2 = c1.submit(inc, 1, key=x.key)
+        y2 = c2.submit(inc, 2, key=y.key)
 
-    assert await x2 == 2
-    assert await y2 == 3
+        assert x2._generation != x._generation
+        assert y2._generation != y._generation
 
-    del x2, y2
-    await async_wait_for(lambda: not s.tasks, timeout=5)
+        assert await x2 == 2
+        assert await y2 == 3
 
-    await c1.close()
-    await c2.close()
+        del x2, y2
+        await async_wait_for(lambda: not s.tasks, timeout=5)
 
 
 @gen_cluster(Worker=Nanny, timeout=60)
diff --git a/distributed/tests/test_publish.py b/distributed/tests/test_publish.py
index 28821bd8aed..ba0b1823b34 100644
--- a/distributed/tests/test_publish.py
+++ b/distributed/tests/test_publish.py
@@ -15,28 +15,25 @@
 
 @gen_cluster()
 async def test_publish_simple(s, a, b):
-    c = Client(s.address, asynchronous=True)
-    f = Client(s.address, asynchronous=True)
-    await asyncio.gather(c, f)
-
-    data = await c.scatter(range(3))
-    await c.publish_dataset(data=data)
-    assert "data" in s.extensions["publish"].datasets
-    assert isinstance(s.extensions["publish"].datasets["data"]["data"], Serialized)
-
-    with pytest.raises(KeyError) as exc_info:
+    async with Client(s.address, asynchronous=True) as c, Client(
+        s.address, asynchronous=True
+    ) as f:
+        data = await c.scatter(range(3))
         await c.publish_dataset(data=data)
+        assert "data" in s.extensions["publish"].datasets
+        assert isinstance(s.extensions["publish"].datasets["data"]["data"], Serialized)
 
-    assert "exists" in str(exc_info.value)
-    assert "data" in str(exc_info.value)
+        with pytest.raises(KeyError) as exc_info:
+            await c.publish_dataset(data=data)
 
-    result = await c.scheduler.publish_list()
-    assert result == ("data",)
+        assert "exists" in str(exc_info.value)
+        assert "data" in str(exc_info.value)
 
-    result = await f.scheduler.publish_list()
-    assert result == ("data",)
+        result = await c.scheduler.publish_list()
+        assert result == ("data",)
 
-    await asyncio.gather(c.close(), f.close())
+        result = await f.scheduler.publish_list()
+        assert result == ("data",)
 
 
 @gen_cluster()
@@ -56,29 +53,27 @@ async def test_publish_non_string_key(s, a, b):
 
 @gen_cluster()
 async def test_publish_roundtrip(s, a, b):
-    c = await Client(s.address, asynchronous=True)
-    f = await Client(s.address, asynchronous=True)
+    async with Client(s.address, asynchronous=True) as c, Client(
+        s.address, asynchronous=True
+    ) as f:
 
-    data = await c.scatter([0, 1, 2])
-    await c.publish_dataset(data=data)
+        data = await c.scatter([0, 1, 2])
+        await c.publish_dataset(data=data)
 
-    assert any(
-        cs.client_key == "published-data" for cs in s.tasks[data[0].key].who_wants
-    )
-    result = await f.get_dataset(name="data")
+        assert any(
+            cs.client_key == "published-data" for cs in s.tasks[data[0].key].who_wants
+        )
+        result = await f.get_dataset(name="data")
 
-    assert len(result) == len(data)
-    out = await f.gather(result)
-    assert out == [0, 1, 2]
+        assert len(result) == len(data)
+        out = await f.gather(result)
+        assert out == [0, 1, 2]
 
-    with pytest.raises(KeyError) as exc_info:
-        await f.get_dataset(name="nonexistent")
-
-    assert "not found" in str(exc_info.value)
-    assert "nonexistent" in str(exc_info.value)
+        with pytest.raises(KeyError) as exc_info:
+            await f.get_dataset(name="nonexistent")
 
-    await c.close()
-    await f.close()
+        assert "not found" in str(exc_info.value)
+        assert "nonexistent" in str(exc_info.value)
 
 
 @gen_cluster(client=True)
@@ -154,29 +149,30 @@ def test_unpublish_multiple_datasets_sync(client):
 @gen_cluster()
 async def test_publish_bag(s, a, b):
     db = pytest.importorskip("dask.bag")
-    c = await Client(s.address, asynchronous=True)
-    f = await Client(s.address, asynchronous=True)
+    async with Client(s.address, asynchronous=True) as c, Client(
+        s.address, asynchronous=True
+    ) as f:
 
-    bag = db.from_sequence([0, 1, 2])
-    bagp = c.persist(bag)
+        bag = db.from_sequence([0, 1, 2])
+        bagp = c.persist(bag)
 
-    assert len(futures_of(bagp)) == 3
-    keys = {f.key for f in futures_of(bagp)}
-    assert keys == set(bag.dask)
+        assert len(futures_of(bagp)) == 3
+        keys = {f.key for f in futures_of(bagp)}
+        assert keys == set(bag.dask)
 
-    await c.publish_dataset(data=bagp)
+        await c.publish_dataset(data=bagp)
 
-    # check that serialization didn't affect original bag's dask
-    assert len(futures_of(bagp)) == 3
+        # check that serialization didn't affect original bag's dask
+        assert len(futures_of(bagp)) == 3
 
-    result = await f.get_dataset("data")
-    assert set(result.dask.keys()) == set(bagp.dask.keys())
-    assert {f.key for f in result.dask.values()} == {f.key for f in bagp.dask.values()}
+        result = await f.get_dataset("data")
+        assert set(result.dask.keys()) == set(bagp.dask.keys())
+        assert {f.key for f in result.dask.values()} == {
+            f.key for f in bagp.dask.values()
+        }
 
-    out = await f.compute(result)
-    assert out == [0, 1, 2]
-    await c.close()
-    await f.close()
+        out = await f.compute(result)
+        assert out == [0, 1, 2]
 
 
 def test_datasets_setitem(client):
diff --git a/distributed/tests/test_queues.py b/distributed/tests/test_queues.py
index 4bcb3bb94f2..9fe47d915ff 100644
--- a/distributed/tests/test_queues.py
+++ b/distributed/tests/test_queues.py
@@ -70,22 +70,20 @@ def test_sync(client):
 
 @gen_cluster()
 async def test_hold_futures(s, a, b):
-    c1 = await Client(s.address, asynchronous=True)
-    future = c1.submit(lambda x: x + 1, 10)
-    q1 = await Queue("q")
-    await q1.put(future)
-    del q1
-    await c1.close()
+    async with Client(s.address, asynchronous=True) as c1:
+        future = c1.submit(lambda x: x + 1, 10)
+        q1 = await Queue("q")
+        await q1.put(future)
+        del q1
 
     await asyncio.sleep(0.1)
 
-    c2 = await Client(s.address, asynchronous=True)
-    q2 = await Queue("q")
-    future2 = await q2.get()
-    result = await future2
+    async with Client(s.address, asynchronous=True) as c1:
+        q2 = await Queue("q")
+        future2 = await q2.get()
+        result = await future2
 
-    assert result == 11
-    await c2.close()
+        assert result == 11
 
 
 @pytest.mark.skip(reason="getting same client from main thread")
@@ -193,31 +191,29 @@ async def test_Future_knows_status_immediately(c, s, a, b):
     q = await Queue("q")
     await q.put(x)
 
-    c2 = await Client(s.address, asynchronous=True)
-    q2 = await Queue("q", client=c2)
-    future = await q2.get()
-    assert future.status == "finished"
-
-    x = c.submit(div, 1, 0)
-    await wait(x)
-    await q.put(x)
+    async with Client(s.address, asynchronous=True) as c2:
+        q2 = await Queue("q", client=c2)
+        future = await q2.get()
+        assert future.status == "finished"
 
-    future2 = await q2.get()
-    assert future2.status == "error"
-    with pytest.raises(Exception):
-        await future2
+        x = c.submit(div, 1, 0)
+        await wait(x)
+        await q.put(x)
 
-    start = time()
-    while True:  # we learn about the true error eventually
-        try:
+        future2 = await q2.get()
+        assert future2.status == "error"
+        with pytest.raises(Exception):
             await future2
-        except ZeroDivisionError:
-            break
-        except Exception:
-            assert time() < start + 5
-            await asyncio.sleep(0.05)
 
-    await c2.close()
+        start = time()
+        while True:  # we learn about the true error eventually
+            try:
+                await future2
+            except ZeroDivisionError:
+                break
+            except Exception:
+                assert time() < start + 5
+                await asyncio.sleep(0.05)
 
 
 @gen_cluster(client=True)
diff --git a/distributed/tests/test_scheduler.py b/distributed/tests/test_scheduler.py
index d75b2aac686..44237027b10 100644
--- a/distributed/tests/test_scheduler.py
+++ b/distributed/tests/test_scheduler.py
@@ -1416,14 +1416,13 @@ async def test_fifo_submission(c, s, w):
 @gen_test()
 async def test_scheduler_file():
     with tmpfile() as fn:
-        s = await Scheduler(scheduler_file=fn, dashboard_address=":0")
-        with open(fn) as f:
-            data = json.load(f)
-        assert data["address"] == s.address
+        async with Scheduler(scheduler_file=fn, dashboard_address=":0") as s:
+            with open(fn) as f:
+                data = json.load(f)
+            assert data["address"] == s.address
 
-        c = await Client(scheduler_file=fn, loop=s.loop, asynchronous=True)
-        await c.close()
-        await s.close()
+            async with Client(scheduler_file=fn, loop=s.loop, asynchronous=True):
+                pass
 
 
 @pytest.mark.parametrize(
diff --git a/distributed/tests/test_utils_test.py b/distributed/tests/test_utils_test.py
index 006627bbbe9..08529aa54b8 100755
--- a/distributed/tests/test_utils_test.py
+++ b/distributed/tests/test_utils_test.py
@@ -273,7 +273,11 @@ def test_new_config():
 def test_lingering_client():
     @gen_cluster()
     async def f(s, a, b):
-        await Client(s.address, asynchronous=True)
+        with pytest.warns(
+            DeprecationWarning,
+            match=r"await Client\(\) is deprecated, use async with Client\(\)",
+        ):
+            await Client(s.address, asynchronous=True)
 
     f()
 
@@ -282,6 +286,7 @@ async def f(s, a, b):
 
 
 def test_lingering_client_2(loop):
+    # assert where the client went
     with cluster() as (s, [a, b]):
         client = Client(s["address"], loop=loop)
 
diff --git a/distributed/tests/test_variable.py b/distributed/tests/test_variable.py
index 3ddf6983b5f..8d963c197ac 100644
--- a/distributed/tests/test_variable.py
+++ b/distributed/tests/test_variable.py
@@ -95,22 +95,20 @@ def test_sync(client):
 
 @gen_cluster()
 async def test_hold_futures(s, a, b):
-    c1 = await Client(s.address, asynchronous=True)
-    future = c1.submit(lambda x: x + 1, 10)
-    x1 = Variable("x")
-    await x1.set(future)
-    del x1
-    await c1.close()
+    async with Client(s.address, asynchronous=True) as c1:
+        future = c1.submit(lambda x: x + 1, 10)
+        x1 = Variable("x")
+        await x1.set(future)
+        del x1
 
     await asyncio.sleep(0.1)
 
-    c2 = await Client(s.address, asynchronous=True)
-    x2 = Variable("x")
-    future2 = await x2.get()
-    result = await future2
+    async with Client(s.address, asynchronous=True) as c2:
+        x2 = Variable("x")
+        future2 = await x2.get()
+        result = await future2
 
-    assert result == 11
-    await c2.close()
+        assert result == 11
 
 
 @gen_cluster(client=True)
@@ -228,31 +226,29 @@ async def test_Future_knows_status_immediately(c, s, a, b):
     v = Variable("x")
     await v.set(x)
 
-    c2 = await Client(s.address, asynchronous=True)
-    v2 = Variable("x", client=c2)
-    future = await v2.get()
-    assert future.status == "finished"
-
-    x = c.submit(div, 1, 0)
-    await wait(x)
-    await v.set(x)
+    async with Client(s.address, asynchronous=True) as c2:
+        v2 = Variable("x", client=c2)
+        future = await v2.get()
+        assert future.status == "finished"
 
-    future2 = await v2.get()
-    assert future2.status == "error"
-    with pytest.raises(Exception):
-        await future2
+        x = c.submit(div, 1, 0)
+        await wait(x)
+        await v.set(x)
 
-    start = time()
-    while True:  # we learn about the true error eventually
-        try:
+        future2 = await v2.get()
+        assert future2.status == "error"
+        with pytest.raises(Exception):
             await future2
-        except ZeroDivisionError:
-            break
-        except Exception:
-            assert time() < start + 5
-            await asyncio.sleep(0.05)
 
-    await c2.close()
+        start = time()
+        while True:  # we learn about the true error eventually
+            try:
+                await future2
+            except ZeroDivisionError:
+                break
+            except Exception:
+                assert time() < start + 5
+                await asyncio.sleep(0.05)
 
 
 @gen_cluster(client=True)
diff --git a/distributed/tests/test_worker.py b/distributed/tests/test_worker.py
index 7e47792ba41..49c4cdeb447 100644
--- a/distributed/tests/test_worker.py
+++ b/distributed/tests/test_worker.py
@@ -1019,7 +1019,10 @@ def f(x):
 @gen_cluster(client=True)
 async def test_get_client_coroutine(c, s, a, b):
     async def f():
-        client = await get_client()
+        # TODO: `await get_client()` raises a deprecation warning to use
+        # `async with get_client()` and that will kill the workers' client
+        # if you do that. We really don't want users to do that.
+        client = get_client()
         future = client.submit(inc, 10)
         result = await future
         return result
diff --git a/distributed/utils_test.py b/distributed/utils_test.py
index e126eb995c1..baa1b6751d2 100644
--- a/distributed/utils_test.py
+++ b/distributed/utils_test.py
@@ -926,11 +926,24 @@ def _(func):
         @config_for_cluster_tests(**{"distributed.comm.timeouts.connect": "5s"})
         @clean(**clean_kwargs)
         def test_func(*outer_args, **kwargs):
-            async def async_fn():
-                result = None
-                with dask.config.set(config):
-                    workers = []
-                    s = False
+            @contextlib.asynccontextmanager
+            async def _client_factory(s):
+                if client:
+                    async with Client(
+                        s.address,
+                        security=security,
+                        asynchronous=True,
+                        **client_kwargs,
+                    ) as c:
+                        yield c
+                else:
+                    yield
+
+            @contextlib.asynccontextmanager
+            async def _cluster_factory():
+                workers = []
+                s = False
+                try:
 
                     for _ in range(60):
                         try:
@@ -951,81 +964,79 @@ async def async_fn():
                             await asyncio.sleep(1)
                         else:
                             workers[:] = ws
-                            args = [s] + workers
                             break
                     if s is False:
                         raise Exception("Could not start cluster")
-                    if client:
-                        c = await Client(
-                            s.address,
-                            security=security,
-                            asynchronous=True,
-                            **client_kwargs,
-                        )
-                        args = [c] + args
-
-                    try:
-                        coro = func(*args, *outer_args, **kwargs)
-                        task = asyncio.create_task(coro)
-                        coro2 = asyncio.wait_for(asyncio.shield(task), timeout)
-                        result = await coro2
-                        validate_state(s, *workers)
-
-                    except asyncio.TimeoutError:
-                        assert task
-                        buffer = io.StringIO()
-                        # This stack indicates where the coro/test is suspended
-                        task.print_stack(file=buffer)
-
-                        if cluster_dump_directory:
-                            await dump_cluster_state(
-                                s,
-                                ws,
-                                output_dir=cluster_dump_directory,
-                                func_name=func.__name__,
-                            )
-
-                        task.cancel()
-                        while not task.cancelled():
-                            await asyncio.sleep(0.01)
-
-                        # Hopefully, the hang has been caused by inconsistent
-                        # state, which should be much more meaningful than the
-                        # timeout
-                        validate_state(s, *workers)
-
-                        # Remove as much of the traceback as possible; it's
-                        # uninteresting boilerplate from utils_test and asyncio
-                        # and not from the code being tested.
-                        raise asyncio.TimeoutError(
-                            f"Test timeout after {timeout}s.\n"
-                            "========== Test stack trace starts here ==========\n"
-                            f"{buffer.getvalue()}"
-                        ) from None
-
-                    except pytest.xfail.Exception:
-                        raise
-
-                    except Exception:
-                        if cluster_dump_directory and not has_pytestmark(
-                            test_func, "xfail"
-                        ):
-                            await dump_cluster_state(
-                                s,
-                                ws,
-                                output_dir=cluster_dump_directory,
-                                func_name=func.__name__,
-                            )
-                        raise
+                    yield s, workers
+                finally:
+                    await end_cluster(s, workers)
+                    await asyncio.wait_for(cleanup_global_workers(), 1)
 
-                    finally:
-                        if client and c.status not in ("closing", "closed"):
-                            await c._close(fast=s.status == Status.closed)
-                        await end_cluster(s, workers)
-                        await asyncio.wait_for(cleanup_global_workers(), 1)
+            async def async_fn():
+                result = None
+                with dask.config.set(config):
+                    async with _cluster_factory() as (s, workers), _client_factory(
+                        s
+                    ) as c:
+                        args = [s] + workers
+                        if c is not None:
+                            args = [c] + args
+                        try:
+                            coro = func(*args, *outer_args, **kwargs)
+                            task = asyncio.create_task(coro)
+                            coro2 = asyncio.wait_for(asyncio.shield(task), timeout)
+                            result = await coro2
+                            validate_state(s, *workers)
+
+                        except asyncio.TimeoutError:
+                            assert task
+                            buffer = io.StringIO()
+                            # This stack indicates where the coro/test is suspended
+                            task.print_stack(file=buffer)
+
+                            if cluster_dump_directory:
+                                await dump_cluster_state(
+                                    s=s,
+                                    ws=workers,
+                                    output_dir=cluster_dump_directory,
+                                    func_name=func.__name__,
+                                )
+
+                            task.cancel()
+                            while not task.cancelled():
+                                await asyncio.sleep(0.01)
+
+                            # Hopefully, the hang has been caused by inconsistent
+                            # state, which should be much more meaningful than the
+                            # timeout
+                            validate_state(s, *workers)
+
+                            # Remove as much of the traceback as possible; it's
+                            # uninteresting boilerplate from utils_test and asyncio
+                            # and not from the code being tested.
+                            raise asyncio.TimeoutError(
+                                f"Test timeout after {timeout}s.\n"
+                                "========== Test stack trace starts here ==========\n"
+                                f"{buffer.getvalue()}"
+                            ) from None
+
+                        except pytest.xfail.Exception:
+                            raise
+
+                        except Exception:
+                            if cluster_dump_directory and not has_pytestmark(
+                                test_func, "xfail"
+                            ):
+                                await dump_cluster_state(
+                                    s=s,
+                                    ws=workers,
+                                    output_dir=cluster_dump_directory,
+                                    func_name=func.__name__,
+                                )
+                            raise
 
                     try:
-                        c = await default_client()
+                        c = default_client()
                     except ValueError:
                         pass
                     else:

From b46814a144d819b17a434c88f30acfc644f23a43 Mon Sep 17 00:00:00 2001
From: Thomas Grainger <tagrain@gmail.com>
Date: Fri, 19 Aug 2022 13:13:05 +0100
Subject: [PATCH 2/3] lingering clients are no-longer deprecated

---
 distributed/tests/test_utils_test.py | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/distributed/tests/test_utils_test.py b/distributed/tests/test_utils_test.py
index 08529aa54b8..f3fd449e16b 100755
--- a/distributed/tests/test_utils_test.py
+++ b/distributed/tests/test_utils_test.py
@@ -273,11 +273,9 @@ def test_new_config():
 def test_lingering_client():
     @gen_cluster()
     async def f(s, a, b):
-        with pytest.warns(
-            DeprecationWarning,
-            match=r"await Client\(\) is deprecated, use async with Client\(\)",
-        ):
-            await Client(s.address, asynchronous=True)
+        # TODO: force async-with here?
+        # see https://github.com/dask/distributed/issues/6616
+        await Client(s.address, asynchronous=True)
 
     f()
 

From a10d2b793b2d3b167f66c020eb7f6ac9cf5e61e2 Mon Sep 17 00:00:00 2001
From: Thomas Grainger <tagrain@gmail.com>
Date: Fri, 19 Aug 2022 13:24:30 +0100
Subject: [PATCH 3/3] Apply suggestions from code review

---
 distributed/tests/test_utils_test.py | 2 +-
 distributed/tests/test_worker.py     | 5 +++--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/distributed/tests/test_utils_test.py b/distributed/tests/test_utils_test.py
index f3fd449e16b..2147982f750 100755
--- a/distributed/tests/test_utils_test.py
+++ b/distributed/tests/test_utils_test.py
@@ -284,7 +284,7 @@ async def f(s, a, b):
 
 
 def test_lingering_client_2(loop):
-    # assert where the client went
+    # TODO: assert where the client went
     with cluster() as (s, [a, b]):
         client = Client(s["address"], loop=loop)
 
diff --git a/distributed/tests/test_worker.py b/distributed/tests/test_worker.py
index 49c4cdeb447..572439ed992 100644
--- a/distributed/tests/test_worker.py
+++ b/distributed/tests/test_worker.py
@@ -1019,9 +1019,10 @@ def f(x):
 @gen_cluster(client=True)
 async def test_get_client_coroutine(c, s, a, b):
     async def f():
-        # TODO: `await get_client()` raises a deprecation warning to use
-        # `async with get_client()` and that will kill the workers' client
+        # TODO: the existence of `await get_client()` implies the possibility
+        # of `async with get_client()` and that will kill the workers' client
         # if you do that. We really don't want users to do that.
+        # https://github.com/dask/distributed/pull/6921/
         client = get_client()
         future = client.submit(inc, 10)
         result = await future