Improve the stack traces given for timeouts with @gen_test.

diff --git a/tornado/test/testing_test.py b/tornado/test/testing_test.py
index 569f32fc43e3a1491ebe49e0d3a5f69947413c47..aabdaced723179626825d1fdda21a482fca294a5 100644 (file)
--- a/tornado/test/testing_test.py
+++ b/tornado/test/testing_test.py
@@ -8,6 +8,7 @@ from tornado.test.util import unittest
 
 import contextlib
 import os
+import traceback
 
 
 @contextlib.contextmanager
@@ -148,8 +149,17 @@ class GenTest(AsyncTestCase):
         def test(self):
             yield gen.Task(self.io_loop.add_timeout, self.io_loop.time() + 1)
 
-        with self.assertRaises(ioloop.TimeoutError):
+        # This can't use assertRaises because we need to inspect the
+        # exc_info triple (and not just the exception object)
+        try:
             test(self)
+            self.fail("did not get expected exception")
+        except ioloop.TimeoutError:
+            # The stack trace should blame the add_timeout line, not just
+            # unrelated IOLoop/testing internals.
+            self.assertIn(
+                "gen.Task(self.io_loop.add_timeout, self.io_loop.time() + 1)",
+                traceback.format_exc())
 
         self.finished = True
 

diff --git a/tornado/testing.py b/tornado/testing.py
index c563bd711c4c69d7476d4d69f264aa7842070b8c..96fdd32b0f4877483c8b24b94176a16a35489f0a 100644 (file)
--- a/tornado/testing.py
+++ b/tornado/testing.py
@@ -17,7 +17,7 @@ try:
     from tornado.httpclient import AsyncHTTPClient
     from tornado.httpserver import HTTPServer
     from tornado.simple_httpclient import SimpleAsyncHTTPClient
-    from tornado.ioloop import IOLoop
+    from tornado.ioloop import IOLoop, TimeoutError
     from tornado import netutil
 except ImportError:
     # These modules are not importable on app engine.  Parts of this module
@@ -455,13 +455,40 @@ def gen_test(func=None, timeout=None):
         timeout = get_async_test_timeout()
 
     def wrap(f):
-        f = gen.coroutine(f)
-
+        # Stack up several decorators to allow us to access the generator
+        # object itself.  In the innermost wrapper, we capture the generator
+        # and save it in an attribute of self.  Next, we run the wrapped
+        # function through @gen.coroutine.  Finally, the coroutine is
+        # wrapped again to make it synchronous with run_sync.
+        #
+        # This is a good case study arguing for either some sort of
+        # extensibility in the gen decorators or cancellation support.
         @functools.wraps(f)
-        def wrapper(self):
-            return self.io_loop.run_sync(
-                functools.partial(f, self), timeout=timeout)
-        return wrapper
+        def pre_coroutine(self):
+            result = f(self)
+            if isinstance(result, types.GeneratorType):
+                self._test_generator = result
+            else:
+                self._test_generator = None
+            return result
+
+        coro = gen.coroutine(pre_coroutine)
+
+        @functools.wraps(coro)
+        def post_coroutine(self):
+            try:
+                return self.io_loop.run_sync(
+                    functools.partial(coro, self), timeout=timeout)
+            except TimeoutError as e:
+                # run_sync raises an error with an unhelpful traceback.
+                # If we throw it back into the generator the stack trace
+                # will be replaced by the point where the test is stopped.
+                self._test_generator.throw(e)
+                # In case the test contains an overly broad except clause,
+                # we may get back here.  In this case re-raise the original
+                # exception, which is better than nothing.
+                raise
+        return post_coroutine
 
     if func is not None:
         # Used like: