diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py index 3121dcb6bdd..14ea3c7f886 100644 --- a/distributed/tests/test_client.py +++ b/distributed/tests/test_client.py @@ -89,6 +89,7 @@ _UnhashableCallable, async_wait_for, asyncinc, + block_on_event, captured_logger, cluster, dec, @@ -727,8 +728,9 @@ async def test_wait(c, s, a, b): @gen_cluster(client=True) async def test_wait_first_completed(c, s, a, b): - x = c.submit(slowinc, 1) - y = c.submit(slowinc, 1) + event = Event() + x = c.submit(block_on_event, event) + y = c.submit(block_on_event, event) z = c.submit(inc, 2) done, not_done = await wait([x, y, z], return_when="FIRST_COMPLETED") @@ -738,6 +740,7 @@ async def test_wait_first_completed(c, s, a, b): assert z.status == "finished" assert x.status == "pending" assert y.status == "pending" + await event.set() @gen_cluster(client=True) diff --git a/distributed/utils_test.py b/distributed/utils_test.py index e851bb0a9b8..031b0e4acc1 100644 --- a/distributed/utils_test.py +++ b/distributed/utils_test.py @@ -38,7 +38,7 @@ import dask -from distributed import Scheduler, system +from distributed import Event, Scheduler, system from distributed import versions as version_module from distributed.batched import BatchedSend from distributed.client import Client, _global_clients, default_client @@ -289,6 +289,10 @@ def lock_inc(x, lock): return x + 1 +def block_on_event(event: Event) -> None: + event.wait() + + class _UnhashableCallable: # FIXME https://github.com/python/mypy/issues/4266 __hash__ = None # type: ignore