async def and yield->await

distributed/cli/tests/test_dask_scheduler.py

@@ -28,9 +28,9 @@
 def test_defaults(loop):
     with popen(["dask-scheduler", "--no-dashboard"]) as proc:
 
-        def f():
+        async def f():
             # Default behaviour is to listen on all addresses
-            yield assert_can_connect_from_everywhere_4_6(8786, timeout=5.0)
+            await assert_can_connect_from_everywhere_4_6(8786, timeout=5.0)
 
         with Client("127.0.0.1:%d" % Scheduler.default_port, loop=loop) as c:
             c.sync(f)
@@ -45,9 +45,9 @@ def f():
 def test_hostport(loop):
     with popen(["dask-scheduler", "--no-dashboard", "--host", "127.0.0.1:8978"]):
 
-        def f():
+        async def f():
             # The scheduler's main port can't be contacted from the outside
-            yield assert_can_connect_locally_4(8978, timeout=5.0)
+            await assert_can_connect_locally_4(8978, timeout=5.0)
 
         with Client("127.0.0.1:8978", loop=loop) as c:
             assert len(c.nthreads()) == 0

distributed/cli/utils.py

@@ -50,10 +50,10 @@ def install_signal_handlers(loop=None, cleanup=None):
     old_handlers = {}
 
     def handle_signal(sig, frame):
-        def cleanup_and_stop():
+        async def cleanup_and_stop():
             try:
                 if cleanup is not None:
-                    yield cleanup(sig)
+                    await cleanup(sig)
             finally:
                 loop.stop()
 

distributed/comm/tests/test_ucx.py

@@ -151,13 +151,13 @@ async def test_ping_pong_data():
 
 
 @gen_test()
-def test_ucx_deserialize():
+async def test_ucx_deserialize():
     # Note we see this error on some systems with this test:
     # `socket.gaierror: [Errno -5] No address associated with hostname`
     # This may be due to a system configuration issue.
     from .test_comms import check_deserialize
 
-    yield check_deserialize("tcp://")
+    await check_deserialize("tcp://")
 
 
 @pytest.mark.asyncio

distributed/dashboard/tests/test_components.py

@@ -21,16 +21,16 @@ def test_basic(Component):
 
 
 @gen_cluster(client=True, clean_kwargs={"threads": False})
-def test_profile_plot(c, s, a, b):
+async def test_profile_plot(c, s, a, b):
     p = ProfilePlot()
     assert not p.source.data["left"]
-    yield c.gather(c.map(slowinc, range(10), delay=0.05))
+    await c.gather(c.map(slowinc, range(10), delay=0.05))
     p.update(a.profile_recent)
     assert len(p.source.data["left"]) >= 1
 
 
 @gen_cluster(client=True, clean_kwargs={"threads": False})
-def test_profile_time_plot(c, s, a, b):
+async def test_profile_time_plot(c, s, a, b):
     from bokeh.io import curdoc
 
     sp = ProfileTimePlot(s, doc=curdoc())
@@ -42,7 +42,7 @@ def test_profile_time_plot(c, s, a, b):
     assert len(sp.source.data["left"]) <= 1
     assert len(ap.source.data["left"]) <= 1
 
-    yield c.gather(c.map(slowinc, range(10), delay=0.05))
+    await c.gather(c.map(slowinc, range(10), delay=0.05))
     ap.trigger_update()
     sp.trigger_update()
-    yield gen.sleep(0.05)
+    await gen.sleep(0.05)