test ActorFuture.result timeout kwarg

Author: graingert

distributed/tests/test_actor.py

@@ -16,6 +16,7 @@
     get_client,
     wait,
 )
+from distributed.actor import _LateLoopEvent
 from distributed.metrics import time
 from distributed.utils_test import cluster, gen_cluster
 
@@ -256,6 +257,27 @@ def test_sync(client):
     assert "distributed.actor" not in repr(future)
 
 
+def test_timeout(client):
+    class Waiter:
+        def __init__(self):
+            self.event = _LateLoopEvent()
+
+        async def set(self):
+            self.event.set()
+
+        async def wait(self):
+            return await self.event.wait()
+
+    event = client.submit(Waiter, actor=True).result()
+    future = event.wait()
+
+    with pytest.raises(asyncio.TimeoutError):
+        future.result(timeout="0.001s")
+
+    event.set().result()
+    assert future.result() is True
+
+
 @gen_cluster(client=True, config={"distributed.comm.timeouts.connect": "1s"})
 async def test_failed_worker(c, s, a, b):
     future = c.submit(Counter, actor=True, workers=[a.address])
@@ -528,11 +550,9 @@ def sleep(self, time):
 
 @gen_cluster(client=True)
 async def test_waiter(c, s, a, b):
-    from tornado.locks import Event
-
     class Waiter:
         def __init__(self):
-            self.event = Event()
+            self.event = _LateLoopEvent()
 
         async def set(self):
             self.event.set()