link_future() now accepts coroutine objects

4 files changed, 48 insertions, 15 deletions

diff --git a/aiogreen.py b/aiogreen.py
index 341e67d..f1ec37c 100644
--- a/aiogreen.py
+++ b/aiogreen.py
@@ -300,15 +300,15 @@ def wrap_greenthread(gt, loop=None):
     return fut
 
 
-def link_future(future):
-    """Wait for a future.
+def link_future(future, loop=None):
+    """Wait for a future, a task, or a coroutine from a greenthread.
 
-    Wait for a future or a task from a greenthread.
     Return the result or raise the exception of the future.
 
     The function must not be called from the greenthread
     of the aiogreen event loop.
     """
+    future = asyncio.async(future, loop=loop)
     if future._loop._greenthread == eventlet.getcurrent():
         raise RuntimeError("link_future() must not be called from "
                            "the greenthread of the aiogreen event loop")
diff --git a/doc/changelog.rst b/doc/changelog.rst
index 8e211f0..7c98919 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -7,6 +7,7 @@ Version 0.3 (development version)
 * :func:`wrap_greenthread` now raises an exception if the greenthread is
   running or already finished. In debug mode, the exception is not more logged
   to sys.stderr for greenthreads.
+* :func:`link_future` now accepts coroutine objects.
 * :func:`link_future` now raises an exception if it is called from the
   greenthread of the aiogreen event loop.
 * Fix eventlet detection of blocking tasks: cancel the alarm when the aiogreen
diff --git a/doc/using.rst b/doc/using.rst
index 2db3979..f9833b8 100644
--- a/doc/using.rst
+++ b/doc/using.rst
@@ -90,9 +90,9 @@ aiogreen specific functions:
 link_future
 -----------
 
-.. function:: link_future(future)
+.. function:: link_future(future, loop=None)
 
-   Wait for a future, or a task, from a greenthread.
+   Wait for a future, a task, or a coroutine from a greenthread.
 
    Return the result or raise the exception of the future.
 
@@ -101,8 +101,9 @@ link_future
 
    .. versionchanged:: 0.3
 
-     An exception is now raised if it is called from the greenthread of the
-     aiogreen event loop.
+      Coroutine objects are also accepted. Added the *loop* parameter.
+      An exception is raised if it is called from the greenthread of the
+      aiogreen event loop.
 
    Example of greenthread waiting for a trollius task. The ``progress()``
    callback is called regulary to see that the event loop in not blocked::
diff --git a/tests/test_eventlet.py b/tests/test_eventlet.py
index 271fdaa..ede6440 100644
--- a/tests/test_eventlet.py
+++ b/tests/test_eventlet.py
@@ -98,17 +98,14 @@ except ImportError:
 
 def greenthread_link_future(result, loop):
     try:
-        t1 = asyncio.async(coro_slow_append(result, 1, 0.2), loop=loop)
-        value = aiogreen.link_future(t1)
+        value = aiogreen.link_future(coro_slow_append(result, 1, 0.020))
         result.append(value)
 
-        t2 = asyncio.async(coro_slow_append(result, 2, 0.1), loop=loop)
-        value = aiogreen.link_future(t2)
+        value = aiogreen.link_future(coro_slow_append(result, 2, 0.010))
         result.append(value)
 
-        t3 = asyncio.async(coro_slow_error(0.001), loop=loop)
         try:
-            value = aiogreen.link_future(t3)
+            value = aiogreen.link_future(coro_slow_error())
         except ValueError as exc:
             result.append(str(exc))
 
@@ -187,10 +184,21 @@ class LinkFutureTests(tests.TestCase):
         self.loop.run_forever()
         self.assertEqual(result, [1, 10, 2, 20, 'error', 4])
 
-    def test_link_future_not_running(self):
+    def test_link_coro(self):
         result = []
+
+        def func(fut):
+            value = aiogreen.link_future(coro_slow_append(result, 3))
+            result.append(value)
+            self.loop.stop()
+
         fut = asyncio.Future(loop=self.loop)
-        event = eventlet.event.Event()
+        eventlet.spawn(func, fut)
+        self.loop.run_forever()
+        self.assertEqual(result, [3, 30])
+
+    def test_link_future_not_running(self):
+        result = []
 
         def func(event, fut):
             event.send('link')
@@ -198,6 +206,8 @@ class LinkFutureTests(tests.TestCase):
             result.append(value)
             self.loop.stop()
 
+        event = eventlet.event.Event()
+        fut = asyncio.Future(loop=self.loop)
         eventlet.spawn(func, event, fut)
         event.wait()
 
@@ -223,6 +233,27 @@ class LinkFutureTests(tests.TestCase):
         self.loop.run_forever()
         self.assertEqual(result, ['error'])
 
+    def test_link_future_wrong_loop(self):
+        result = []
+        loop2 = asyncio.new_event_loop()
+        self.addCleanup(loop2.close)
+
+        def func(fut):
+            try:
+                value = aiogreen.link_future(fut, loop=loop2)
+            except Exception as exc:
+                result.append(str(exc))
+            else:
+                result.append(value)
+            self.loop.stop()
+
+        fut = asyncio.Future(loop=self.loop)
+        self.loop.call_soon(func, fut)
+        self.loop.call_soon(fut.set_result, 'unused')
+        self.loop.run_forever()
+        self.assertEqual(result[0],
+                         'loop argument must agree with Future')
+
 
 class WrapGreenthreadTests(tests.TestCase):
     def test_wrap_greenthread(self):