Add the TestCase.with_timeout decorator.

asynctest/case.py

@@ -44,6 +44,7 @@
 
 import asyncio
 import functools
+import os
 import types
 import unittest.case
 import warnings
@@ -183,6 +184,35 @@ def test_something(self):
     #: Event loop created and set as default event loop during the test.
     loop = None
 
+    #: Default timeout for the :meth:`~asynctest.TestCase.with_timeout` decorator.
+    #: The :envvar:`ASYNCTEST_TIMEOUT` environment variable can be used to override default.
+    default_timeout = 1.0
+
+    @classmethod
+    def with_timeout(cls, timeout=None):
+        """
+        Raise :exc:`asyncio.TimeoutError` if decorated coroutine does not finish
+        within specified timeout.
+
+        :param timeout: Timeout for the wrapped coroutine.
+
+        :see: :attr:`~asynctest.TestCase.default_timeout`
+        """
+        def decorator(coro):
+            @functools.wraps(coro)
+            @asyncio.coroutine
+            def wrapper(self, *args, **kwargs):
+                nonlocal timeout
+
+                if timeout is None:
+                    timeout = float(os.environ.get('ASYNCTEST_TIMEOUT', self.default_timeout))
+
+                return (yield from asyncio.wait_for(coro(self, *args, **kwargs), timeout=timeout))
+
+            return wrapper
+
+        return decorator
+
     def _init_loop(self):
         if self.use_default_loop:
             self.loop = asyncio.get_event_loop()

test/test_case.py

@@ -338,6 +338,40 @@ def runTest(self):
                 if outcome:
                     self.assertTrue(outcome.wasSuccessful())
 
+    def test_with_timeout(self):
+        class ArgumentWithTimeoutTest(asynctest.TestCase):
+            @asynctest.TestCase.with_timeout(0.1)
+            @asyncio.coroutine
+            def runTest(self):
+                yield from asyncio.sleep(0.5)
+
+        class OverrideWithTimeoutTest(asynctest.TestCase):
+            default_timeout = 0.1
+
+            @asynctest.TestCase.with_timeout()
+            @asyncio.coroutine
+            def runTest(self):
+                yield from asyncio.sleep(0.5)
+
+        class EnvironWithTimeoutTest(asynctest.TestCase):
+            @asynctest.TestCase.with_timeout()
+            @asyncio.coroutine
+            def runTest(self):
+                yield from asyncio.sleep(0.5)
+
+        cases = [ArgumentWithTimeoutTest(), OverrideWithTimeoutTest()]
+
+        for case in cases:
+            with self.subTest(case=case):
+                with self.assertRaises(asyncio.TimeoutError):
+                    case.debug()
+
+        with unittest.mock.patch.dict(os.environ, {'ASYNCTEST_TIMEOUT': '0.1'}):
+            case = EnvironWithTimeoutTest()
+            with self.subTest(case=case):
+                with self.assertRaises(asyncio.TimeoutError):
+                    case.debug()
+
 
 @unittest.skipIf(sys.platform == "win32", "Tests specific to Unix")
 class Test_TestCase_and_ChildWatcher(_TestCase):