Add test for as_completed for loops in Python 2

distributed/client.py

@@ -45,7 +45,8 @@
 from .utils_comm import (WrappedKey, unpack_remotedata, pack_data,
                          scatter_to_workers, gather_from_workers)
 from .cfexecutor import ClientExecutor
-from .compatibility import Queue as pyQueue, Empty, isqueue, html_escape
+from .compatibility import (Queue as pyQueue, Empty, isqueue, html_escape,
+        StopAsyncIteration)
 from .core import connect, rpc, clean_exception, CommClosedError, PooledRPCCall
 from .metrics import time
 from .node import Node

distributed/compatibility.py

@@ -18,6 +18,8 @@
     PY3 = False
     ConnectionRefusedError = OSError
     FileExistsError = OSError
+    class StopAsyncIteration(Exception):
+        pass
 
     import gzip
 
@@ -71,6 +73,7 @@ def iscoroutinefunction(func):
     from gzip import compress as gzip_compress
     ConnectionRefusedError = ConnectionRefusedError
     FileExistsError = FileExistsError
+    StopAsyncIteration = StopAsyncIteration
 
     def isqueue(o):
         return isinstance(o, Queue)

distributed/tests/test_as_completed.py

@@ -8,10 +8,9 @@
 
 from distributed import Client
 from distributed.client import _as_completed, as_completed, _first_completed
-from distributed.compatibility import Empty
+from distributed.compatibility import Empty, StopAsyncIteration, Queue
 from distributed.utils_test import cluster, gen_cluster, inc
 from distributed.utils_test import loop  # noqa: F401
-from distributed.compatibility import Queue
 
 
 @gen_cluster(client=True)
@@ -152,3 +151,20 @@ def _():
             result = list(ac)
 
             assert result == [x]
+
+
+@gen_cluster(client=True)
+def test_async_for_py2_equivalent(c, s, a, b):
+    futures = c.map(sleep, [0.01] * 3, pure=False)
+    seq = as_completed(futures)
+    x = yield seq.__anext__()
+    y = yield seq.__anext__()
+    z = yield seq.__anext__()
+
+    assert x.done()
+    assert y.done()
+    assert z.done()
+    assert x.key != y.key
+
+    with pytest.raises(StopAsyncIteration):
+        yield seq.__anext__()