wrap_greenthread(): catch missing run attribute

Add 2 tests: no run attribute, wrap exception.

3 files changed, 24 insertions, 1 deletions

diff --git a/aiogreen.py b/aiogreen.py
index dead9a4..012bab3 100644
--- a/aiogreen.py
+++ b/aiogreen.py
@@ -269,6 +269,8 @@ def wrap_greenthread(gt, loop=None):
 
     The greenthread must be wrapped before its execution starts. If the
     greenthread is running or already finished, an exception is raised.
+
+    For greenlets, the run attribute must be set.
     """
     if loop is None:
         loop = asyncio.get_event_loop()
@@ -295,7 +297,11 @@ def wrap_greenthread(gt, loop=None):
                 fut.set_result(result)
         gt.run = wrap_func
     else:
-        orig_func = gt.run
+        try:
+            orig_func = gt.run
+        except AttributeError:
+            raise RuntimeError("wrap_greenthread: the run attribute "
+                               "of the greenlet is not set")
         def wrap_func(*args, **kw):
             try:
                 result = orig_func(*args, **kw)
diff --git a/doc/using.rst b/doc/using.rst
index d36a34a..810af1b 100644
--- a/doc/using.rst
+++ b/doc/using.rst
@@ -197,6 +197,8 @@ wrap_greenthread
    The greenthread must be wrapped before its execution starts.  If the
    greenthread is running or already finished, an exception is raised.
 
+   For greenlets, the ``run`` attribute must be set.
+
    .. versionchanged:: 0.3
 
      An exception is now raised if the greenthread is running or already
diff --git a/tests/test_greenlet.py b/tests/test_greenlet.py
index 67cde61..8ee8a06 100644
--- a/tests/test_greenlet.py
+++ b/tests/test_greenlet.py
@@ -14,6 +14,21 @@ class WrapGreenletTests(tests.TestCase):
         result = self.loop.run_until_complete(fut)
         self.assertEqual(result, 15)
 
+    def test_wrap_greenlet_exc(self):
+        def func():
+            raise ValueError(7)
+
+        gl = greenlet.greenlet(func)
+        fut = aiogreen.wrap_greenthread(gl)
+        gl.switch()
+        self.assertRaises(ValueError, self.loop.run_until_complete, fut)
+
+    def test_wrap_greenlet_no_run_attr(self):
+        gl = greenlet.greenlet()
+        msg = "wrap_greenthread: the run attribute of the greenlet is not set"
+        self.assertRaisesRegexp(RuntimeError, msg,
+                                aiogreen.wrap_greenthread, gl)
+
     def test_wrap_greenlet_running(self):
         def func(value):
             gl = greenlet.getcurrent()