Harden decorators

Fix flaky tests

distributed/client.py

@@ -102,7 +102,6 @@ def _get_global_client():
             return c
         else:
             del _global_clients[k]
-    del L
     return None
 
 

distributed/dashboard/tests/test_scheduler_bokeh.py

@@ -62,7 +62,7 @@ async def test_simple(c, s, a, b):
 
 
 @gen_cluster(client=True, worker_kwargs={"dashboard": True})
-def test_basic(c, s, a, b):
+async def test_basic(c, s, a, b):
     for component in [TaskStream, SystemMonitor, Occupancy, StealingTimeSeries]:
         ss = component(s)
 

distributed/deploy/tests/test_adaptive.py

@@ -215,7 +215,7 @@ async def test_avoid_churn(cleanup):
             assert len(adapt.log) == 1
 
 
-@gen_test(timeout=None)
+@pytest.mark.asyncio
 async def test_adapt_quickly():
     """ We want to avoid creating and deleting workers frequently
 
@@ -332,7 +332,7 @@ def test_basic_no_loop(loop):
             loop.add_callback(loop.stop)
 
 
-@gen_test(timeout=None)
+@pytest.mark.asyncio
 async def test_target_duration():
     """ Ensure that redefining adapt with a lower maximum removes workers """
     with dask.config.set(

distributed/tests/test_client.py

@@ -2293,7 +2293,10 @@ async def test_cancel_collection(c, s, a, b):
     await c.cancel(x)
     await c.cancel([x])
     assert all(f.cancelled() for f in L)
-    assert not s.tasks
+    start = time()
+    while s.tasks:
+        assert time() < start + 1
+        await time.sleep(0.01)
 
 
 def test_cancel(c):

distributed/tests/test_core.py

@@ -732,7 +732,7 @@ async def f():
 
 
 @gen_cluster()
-def test_thread_id(s, a, b):
+async def test_thread_id(s, a, b):
     assert s.thread_id == a.thread_id == b.thread_id == threading.get_ident()
 
 

distributed/tests/test_nanny.py

@@ -28,7 +28,8 @@
 )
 
 
-@gen_cluster(nthreads=[])
+# FIXME why does this leave behind unclosed Comm objects?
+@gen_cluster(nthreads=[], allow_unclosed=True)
 async def test_nanny(s):
     async with Nanny(s.address, nthreads=2, loop=s.loop) as n:
         async with rpc(n.address) as nn:
@@ -68,7 +69,7 @@ async def test_many_kills(s):
 
 
 @gen_cluster(Worker=Nanny)
-def test_str(s, a, b):
+async def test_str(s, a, b):
     assert a.worker_address in str(a)
     assert a.worker_address in repr(a)
     assert str(a.nthreads) in str(a)

distributed/tests/test_scheduler.py

@@ -49,7 +49,7 @@
 
 
 @gen_cluster()
-def test_administration(s, a, b):
+async def test_administration(s, a, b):
     assert isinstance(s.address, str)
     assert s.address in str(s)
     assert str(sum(s.nthreads.values())) in repr(s)
@@ -478,7 +478,7 @@ def test_dumps_task():
 
 
 @gen_cluster()
-def test_ready_remove_worker(s, a, b):
+async def test_ready_remove_worker(s, a, b):
     s.update_graph(
         tasks={"x-%d" % i: dumps_task((inc, i)) for i in range(20)},
         keys=["x-%d" % i for i in range(20)],
@@ -1279,7 +1279,7 @@ async def test_reschedule(c, s, a, b):
 
 
 @gen_cluster(client=True, nthreads=[("127.0.0.1", 1)] * 2)
-def test_reschedule_warns(c, s, a, b):
+async def test_reschedule_warns(c, s, a, b):
     with captured_logger(logging.getLogger("distributed.scheduler")) as sched:
         s.reschedule(key="__this-key-does-not-exist__")
 
@@ -1515,7 +1515,7 @@ def qux(x):
 
 
 @gen_cluster(client=True)
-def test_collect_versions(c, s, a, b):
+async def test_collect_versions(c, s, a, b):
     cs = s.clients[c.id]
     (w1, w2) = s.workers.values()
     assert cs.versions
@@ -1584,7 +1584,7 @@ async def f(dask_worker):
 
 
 @gen_cluster()
-def test_workerstate_clean(s, a, b):
+async def test_workerstate_clean(s, a, b):
     ws = s.workers[a.address].clean()
     assert ws.address == a.address
     b = pickle.dumps(ws)

distributed/tests/test_semaphore.py

@@ -325,17 +325,15 @@ async def test_oversubscribing_leases(c, s, a, b):
     accept new leases as long as the semaphore is oversubscribed.
 
     Oversubscription may occur if tasks hold the GIL for a longer time than the
-    lease-timeout is configured causing the lease refreshs to go stale and
-    timeout.
+    lease-timeout is configured causing the lease refresh to go stale and timeout.
 
     We cannot protect ourselves entirely from this but we can ensure that while
     a task with a timed out lease is still running, we block further
     acquisitions until we return to normal.
 
     An example would be a task which continuously locks the GIL for a longer
-    time than the lease timeout but this continous lock only makes up a
+    time than the lease timeout but this continuous lock only makes up a
     fraction of the tasks runtime.
-
     """
     # GH3705
 

distributed/tests/test_steal.py

@@ -569,7 +569,9 @@ async def assert_balanced(inp, expected, c, s, *workers):
     ],
 )
 def test_balance(inp, expected):
-    test = lambda *args, **kwargs: assert_balanced(inp, expected, *args, **kwargs)
+    async def test(*args, **kwargs):
+        await assert_balanced(inp, expected, *args, **kwargs)
+
     test = gen_cluster(
         client=True,
         nthreads=[("127.0.0.1", 1)] * len(inp),

distributed/tests/test_tls_functional.py

@@ -12,7 +12,7 @@
 
 
 @gen_tls_cluster(client=True)
-def test_basic(c, s, a, b):
+async def test_basic(c, s, a, b):
     pass
 
 

distributed/tests/test_utils_test.py

@@ -1,4 +1,3 @@
-import asyncio
 from contextlib import contextmanager
 import socket
 import threading
@@ -43,7 +42,7 @@ def test_cluster(loop):
 
 
 @gen_cluster(client=True)
-def test_gen_cluster(c, s, a, b):
+async def test_gen_cluster(c, s, a, b):
     assert isinstance(c, Client)
     assert isinstance(s, Scheduler)
     for w in [a, b]:
@@ -68,20 +67,25 @@ async def f(c, s, a, b):
 
 
 @gen_cluster(client=False)
-def test_gen_cluster_without_client(s, a, b):
+async def test_gen_cluster_without_client(s, a, b):
     assert isinstance(s, Scheduler)
     for w in [a, b]:
         assert isinstance(w, Worker)
     assert s.nthreads == {w.address: w.nthreads for w in [a, b]}
 
+    async with Client(s.address, asynchronous=True) as c:
+        future = c.submit(lambda x: x + 1, 1)
+        result = await future
+        assert result == 2
+
 
 @gen_cluster(
     client=True,
     scheduler="tls://127.0.0.1",
     nthreads=[("tls://127.0.0.1", 1), ("tls://127.0.0.1", 2)],
     security=tls_only_security(),
 )
-def test_gen_cluster_tls(e, s, a, b):
+async def test_gen_cluster_tls(e, s, a, b):
     assert isinstance(e, Client)
     assert isinstance(s, Scheduler)
     assert s.address.startswith("tls://")
@@ -91,11 +95,6 @@ def test_gen_cluster_tls(e, s, a, b):
     assert s.nthreads == {w.address: w.nthreads for w in [a, b]}
 
 
-@gen_test()
-async def test_gen_test():
-    await asyncio.sleep(0.01)
-
-
 @contextmanager
 def _listen(delay=0):
     serv = socket.socket()
@@ -177,16 +176,3 @@ def test_tls_cluster(tls_client):
 async def test_tls_scheduler(security, cleanup):
     async with Scheduler(security=security, host="localhost") as s:
         assert s.address.startswith("tls")
-
-
-@gen_cluster()
-async def test_gen_cluster_async(s, a, b):  # flake8: noqa
-    async with Client(s.address, asynchronous=True) as c:
-        future = c.submit(lambda x: x + 1, 1)
-        result = await future
-        assert result == 2
-
-
-@gen_test()
-async def test_gen_test_async():  # flake8: noqa
-    await asyncio.sleep(0.001)

distributed/tests/test_worker.py

@@ -837,7 +837,7 @@ def test_worker_dir(worker):
     with tmpfile() as fn:
 
         @gen_cluster(client=True, worker_kwargs={"local_directory": fn})
-        def test_worker_dir(c, s, a, b):
+        async def test_worker_dir(c, s, a, b):
             directories = [w.local_directory for w in s.workers.values()]
             assert all(d.startswith(fn) for d in directories)
             assert len(set(directories)) == 2  # distinct
@@ -1244,7 +1244,7 @@ async def test_avoid_memory_monitor_if_zero_limit(c, s):
         "distributed.worker.memory.target": False,
     },
 )
-def test_dict_data_if_no_spill_to_disk(s, w):
+async def test_dict_data_if_no_spill_to_disk(s, w):
     assert type(w.data) is dict
 
 

distributed/utils_test.py

@@ -1,5 +1,6 @@
 import asyncio
 import collections
+import gc
 from contextlib import contextmanager
 import copy
 import functools
@@ -33,7 +34,6 @@
 
 import dask
 from tlz import merge, memoize, assoc
-from tornado import gen
 from tornado.ioloop import IOLoop
 
 from . import system
@@ -768,11 +768,9 @@ async def test_foo():
     def _(func):
         def test_func():
             with clean() as loop:
-                if iscoroutinefunction(func):
-                    cor = func
-                else:
-                    cor = gen.coroutine(func)
-                loop.run_sync(cor, timeout=timeout)
+                if not iscoroutinefunction(func):
+                    raise ValueError("@gen_test should wrap async def functions")
+                loop.run_sync(func, timeout=timeout)
 
         return test_func
 
@@ -856,6 +854,7 @@ def gen_cluster(
     active_rpc_timeout=1,
     config={},
     clean_kwargs={},
+    allow_unclosed=False,
 ):
     from distributed import Client
 
@@ -878,10 +877,10 @@ async def test_foo(scheduler, worker1, worker2):
     )
 
     def _(func):
-        if not iscoroutinefunction(func):
-            func = gen.coroutine(func)
-
         def test_func():
+            if not iscoroutinefunction(func):
+                raise ValueError("@gen_cluster should wrap async def functions")
+
             result = None
             workers = []
             with clean(timeout=active_rpc_timeout, **clean_kwargs) as loop:
@@ -905,6 +904,7 @@ async def coro():
                                     "Failed to start gen_cluster, retrying",
                                     exc_info=True,
                                 )
+                                await asyncio.sleep(1)
                             else:
                                 workers[:] = ws
                                 args = [s] + workers
@@ -940,16 +940,28 @@ async def coro():
                         else:
                             await c._close(fast=True)
 
-                        for i in range(5):
-                            if all(c.closed() for c in Comm._instances):
-                                break
-                            else:
+                        def get_unclosed():
+                            return [c for c in Comm._instances if not c.closed()] + [
+                                c
+                                for c in _global_clients.values()
+                                if c.status != "closed"
+                            ]
+
+                        try:
+                            start = time()
+                            while time() < start + 5:
+                                gc.collect()
+                                if not get_unclosed():
+                                    break
                                 await asyncio.sleep(0.05)
-                        else:
-                            L = [c for c in Comm._instances if not c.closed()]
+                            else:
+                                if allow_unclosed:
+                                    print(f"Unclosed Comms: {get_unclosed()}")
+                                else:
+                                    raise RuntimeError("Unclosed Comms", get_unclosed())
+                        finally:
                             Comm._instances.clear()
-                            # raise ValueError("Unclosed Comms", L)
-                            print("Unclosed Comms", L)
+                            _global_clients.clear()
 
                         return result