Make gen_test's timeout configurable

diff --git a/tornado/test/testing_test.py b/tornado/test/testing_test.py
index dc9e0866ed4c56b8b65528d350a64bea48bf5415..40c2fe53aedc47824ee0aa35e46ef89258e0e7ab 100644 (file)
--- a/tornado/test/testing_test.py
+++ b/tornado/test/testing_test.py
@@ -2,10 +2,13 @@
 
 from __future__ import absolute_import, division, print_function, with_statement
 
-from tornado import gen
+from tornado import gen, ioloop
 from tornado.testing import AsyncTestCase, gen_test
 from tornado.test.util import unittest
 
+import functools
+import os
+
 
 class AsyncTestCaseTest(AsyncTestCase):
     def test_exception_in_callback(self):
@@ -74,5 +77,55 @@ class GenTest(AsyncTestCase):
         yield gen.Task(self.io_loop.add_callback)
         self.finished = True
 
+    def test_timeout(self):
+        # Set a short timeout and exceed it.
+        @gen_test(timeout=0.1)
+        def test(self):
+            yield gen.Task(self.io_loop.add_timeout, self.io_loop.time() + 1)
+
+        with self.assertRaises(ioloop.TimeoutError):
+            test(self)
+
+        self.finished = True
+
+    def test_no_timeout(self):
+        # A test that does not exceed its timeout should succeed.
+        @gen_test(timeout=1)
+        def test(self):
+            yield gen.Task(self.io_loop.add_timeout, self.io_loop.time() + 0.1)
+
+        test(self)
+        self.finished = True
+
+    def test_timeout_environment_variable(self):
+        time = self.io_loop.time
+        add_timeout = self.io_loop.add_timeout
+        old_timeout = os.environ.get('TIMEOUT')
+        try:
+            os.environ['TIMEOUT'] = '0.1'
+
+            @gen_test(timeout=0.5)
+            def test_long_timeout(self):
+                yield gen.Task(add_timeout, time() + 0.25)
+
+            # Uses provided timeout of 0.5 seconds, doesn't time out.
+            self.io_loop.run_sync(
+                functools.partial(test_long_timeout, self))
+
+            @gen_test(timeout=0.01)
+            def test_short_timeout(self):
+                yield gen.Task(add_timeout, time() + 1)
+
+            # Uses environment TIMEOUT of 0.1, times out.
+            with self.assertRaises(ioloop.TimeoutError):
+                test_short_timeout(self)
+
+            self.finished = True
+        finally:
+            if old_timeout is None:
+                del os.environ['TIMEOUT']
+            else:
+                os.environ['TIMEOUT'] = old_timeout
+
 if __name__ == '__main__':
     unittest.main()

diff --git a/tornado/testing.py b/tornado/testing.py
index 61c8430611edb028dfc77b61c959179b58248218..1330ea2055ade854c1d591201429fe9eb2846897 100644 (file)
--- a/tornado/testing.py
+++ b/tornado/testing.py
@@ -32,6 +32,7 @@ from tornado.log import gen_log
 from tornado.stack_context import ExceptionStackContext
 from tornado.util import raise_exc_info, basestring_type
 import functools
+import inspect
 import logging
 import os
 import re
@@ -354,7 +355,7 @@ class AsyncHTTPSTestCase(AsyncHTTPTestCase):
         return 'https'
 
 
-def gen_test(f):
+def gen_test(timeout=None):
     """Testing equivalent of ``@gen.coroutine``, to be applied to test methods.
 
     ``@gen.coroutine`` cannot be used on tests because the `.IOLoop` is not
@@ -368,13 +369,49 @@ def gen_test(f):
             def test_something(self):
                 response = yield gen.Task(self.fetch('/'))
 
+    By default, ``@gen_test`` times out after 5 seconds. The timeout may be
+    overridden globally with the TIMEOUT environment variable, or for each
+    test with the ``timeout`` parameter:
+
+        class MyTest(AsyncHTTPTestCase):
+            @gen_test(timeout=10)
+            def test_something_slow(self):
+                response = yield gen.Task(self.fetch('/'))
+
+    If both the environment variable and the parameter are set, ``gen_test``
+    uses the maximum of the two.
     """
-    f = gen.coroutine(f)
+    try:
+        env_timeout = float(os.environ.get('TIMEOUT'))
+    except (ValueError, TypeError):
+        env_timeout = None
+
+    def wrap(f):
+        f = gen.coroutine(f)
+
+        @functools.wraps(f)
+        def wrapper(self):
+            return self.io_loop.run_sync(
+                functools.partial(f, self), timeout=timeout)
+        return wrapper
+
+    if inspect.isfunction(timeout):
+        # Used like:
+        #     @gen_test
+        #     def f(self):
+        #         pass
+        # The 'timeout' parameter is actually the test function.
+        f = timeout
+        timeout = env_timeout or 5
+        return wrap(f)
+    else:
+        # Used like @gen_test(timeout=10) or @gen_test(10).
+        if env_timeout is not None:
+            timeout = max(float(timeout), env_timeout)
+        else:
+            timeout = float(timeout)
 
-    @functools.wraps(f)
-    def wrapper(self):
-        return self.io_loop.run_sync(functools.partial(f, self), timeout=5)
-    return wrapper
+        return wrap
 
 
 # Without this attribute, nosetests will try to run gen_test as a test