From 08c52bf4011268639f341b909792b6b433db81e9 Mon Sep 17 00:00:00 2001
From: Ilya Kulakov <kulakov.ilya@gmail.com>
Date: Mon, 20 Nov 2017 17:50:48 -0800
Subject: [PATCH] Add the TestCase.with_timeout decorator.

---
 asynctest/case.py | 30 ++++++++++++++++++++++++++++++
 test/test_case.py | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+)

diff --git a/asynctest/case.py b/asynctest/case.py
index a0a9151..66284d9 100644
--- a/asynctest/case.py
+++ b/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()
diff --git a/test/test_case.py b/test/test_case.py
index 2df1168..2cdb1bc 100644
--- a/test/test_case.py
+++ b/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):