async... await in tests

.gitignore

@@ -17,3 +17,4 @@ dask-worker-space/
 *.swp
 .ycm_extra_conf.py
 tags
+.ipynb_checkpoints

distributed/cli/tests/test_dask_scheduler.py

@@ -10,7 +10,6 @@
 import tempfile
 from time import sleep
 
-from tornado import gen
 from click.testing import CliRunner
 
 import distributed
@@ -29,12 +28,9 @@
 def test_defaults(loop):
     with popen(["dask-scheduler", "--no-dashboard"]) as proc:
 
-        @gen.coroutine
-        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)
-            ]  # main port
+            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)
@@ -49,12 +45,9 @@ def f():
 def test_hostport(loop):
     with popen(["dask-scheduler", "--no-dashboard", "--host", "127.0.0.1:8978"]):
 
-        @gen.coroutine
-        def f():
-            yield [
-                # The scheduler's main port can't be contacted from the outside
-                assert_can_connect_locally_4(8978, timeout=5.0)
-            ]
+        async def f():
+            # The scheduler's main port can't be contacted from the outside
+            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

@@ -1,4 +1,3 @@
-from tornado import gen
 from tornado.ioloop import IOLoop
 
 
@@ -51,11 +50,10 @@ def install_signal_handlers(loop=None, cleanup=None):
     old_handlers = {}
 
     def handle_signal(sig, frame):
-        @gen.coroutine
-        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/client.py

@@ -102,7 +102,6 @@ def _get_global_client():
             return c
         else:
             del _global_clients[k]
-    del L
     return None
 
 
@@ -1339,6 +1338,10 @@ def close(self, timeout=no_default):
             timeout = self._timeout * 2
         # XXX handling of self.status here is not thread-safe
         if self.status == "closed":
+            if self.asynchronous:
+                future = asyncio.Future()
+                future.set_result(None)
+                return future
             return
         self.status = "closing"
 

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

@@ -1,9 +1,10 @@
+import asyncio
+
 import pytest
 
 pytest.importorskip("bokeh")
 
 from bokeh.models import ColumnDataSource, Model
-from tornado import gen
 
 from distributed.utils_test import slowinc, gen_cluster
 from distributed.dashboard.components.shared import (
@@ -21,16 +22,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.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 +43,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.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 asyncio.sleep(0.05)