Timeout if an individual test takes >10 minutes.

BUG=chromium:323675
TEST=Run with example case where test case takes too long.
TEST=paladin and release trybots

Change-Id: I2ce82484ef15efc9fa17d9a837c43ca7efbb6e52
Reviewed-on: https://chromium-review.googlesource.com/178310
Reviewed-by: David James <davidjames@chromium.org>
Tested-by: David James <davidjames@chromium.org>
Commit-Queue: David James <davidjames@chromium.org>
diff --git a/lib/cros_build_lib.py b/lib/cros_build_lib.py
index b3007e2..55dcb48 100644
--- a/lib/cros_build_lib.py
+++ b/lib/cros_build_lib.py
@@ -1113,16 +1113,33 @@
 
 def TimeoutDecorator(max_time):
   """Decorator used to ensure a func is interrupted if it's running too long."""
-  def decorator(functor):
-    def wrapper(self, *args, **kwds):
-      with Timeout(max_time):
-        return functor(self, *args, **kwds)
+  # Save off the built-in versions of time.time, signal.signal, and
+  # signal.alarm, in case they get mocked out later. We want to ensure that
+  # tests don't accidentally mock out the functions used by Timeout.
+  def _Save():
+    return time.time, signal.signal, signal.alarm
+  def _Restore(values):
+    (time.time, signal.signal, signal.alarm) = values
+  builtins = _Save()
 
-    wrapper.__module__ = functor.__module__
-    wrapper.__name__ = functor.__name__
-    wrapper.__doc__ = functor.__doc__
-    return wrapper
-  return decorator
+  def NestedTimeoutDecorator(func):
+    @functools.wraps(func)
+    def TimeoutWrapper(*args, **kwargs):
+      new = _Save()
+      try:
+        _Restore(builtins)
+        with Timeout(max_time):
+          _Restore(new)
+          try:
+            func(*args, **kwargs)
+          finally:
+            _Restore(builtins)
+      finally:
+        _Restore(new)
+
+    return TimeoutWrapper
+
+  return NestedTimeoutDecorator
 
 
 class ContextManagerStack(object):
diff --git a/lib/cros_test_lib.py b/lib/cros_test_lib.py
index d2592a5..729e0bb 100644
--- a/lib/cros_test_lib.py
+++ b/lib/cros_test_lib.py
@@ -180,21 +180,29 @@
 
 
 class StackedSetup(type):
-  """Metaclass that extracts automatically stacks setUp and tearDown calls.
+  """Metaclass to simplify unit testing and make it more robust.
 
-  Basically this exists to make it easier to do setUp *correctly*, while also
-  suppressing some unittests misbehaviours- for example, the fact that if a
-  setUp throws an exception the corresponding tearDown isn't ran.  This sorts
-  it.
+  A metaclass alters the way that classes are initialized, enabling us to
+  modify the class dictionary prior to the class being created. We use this
+  feature here to modify the way that unit tests work a bit.
 
-  Usage of it is via usual metaclass approach; just set
-  `__metaclass__ = StackedSetup`.
+  This class does three things:
+    1) When a test case is set up or torn down, we now run all setUp and
+       tearDown methods in the inheritance tree.
+    2) If a setUp or tearDown method fails, we still run tearDown methods
+       for any test classes that were partially or completely set up.
+    3) All test cases time out after TEST_CASE_TIMEOUT seconds.
 
-  Note that this metaclass is designed such that because this is a metaclass,
-  rather than just a scope mutator, all derivative classes derive from this
-  metaclass; thus all derivative TestCase classes get automatic stacking.
+  To use this class, set the following in your class:
+    __metaclass__ = StackedSetup
+
+  Since cros_test_lib.TestCase uses this metaclass, all derivatives of TestCase
+  also inherit the above behavior (unless they override the __metaclass__
+  attribute manually.)
   """
 
+  TEST_CASE_TIMEOUT = 10 * 60
+
   def __new__(mcs, name, bases, scope):
     """Generate the new class with pointers to original funcs & our helpers"""
     if 'setUp' in scope:
@@ -205,6 +213,14 @@
       scope['__raw_tearDown__'] = scope.pop('tearDown')
     scope['tearDown'] = mcs._stacked_tearDown
 
+    # Modify all test* methods to time out after TEST_CASE_TIMEOUT seconds.
+    timeout = scope.get('TEST_CASE_TIMEOUT', StackedSetup.TEST_CASE_TIMEOUT)
+    if timeout is not None:
+      for name, func in scope.iteritems():
+        if name.startswith('test') and hasattr(func, '__call__'):
+          wrapper = cros_build_lib.TimeoutDecorator(timeout)
+          scope[name] = wrapper(func)
+
     return type.__new__(mcs, name, bases, scope)
 
   @staticmethod
diff --git a/lib/cros_test_lib_unittest.py b/lib/cros_test_lib_unittest.py
index 93d549d..c5ff28f 100755
--- a/lib/cros_test_lib_unittest.py
+++ b/lib/cros_test_lib_unittest.py
@@ -6,10 +6,13 @@
 
 import os
 import sys
+import time
+import unittest
 
 sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)),
                                 '..', '..'))
 
+from chromite.lib import cros_build_lib
 from chromite.lib import cros_test_lib
 from chromite.lib import cros_build_lib_unittest
 from chromite.lib import partial_mock
@@ -175,5 +178,26 @@
     self.assertEquals(self.Mockable.TO_BE_MOCKED3, 20)
 
 
+class TestCaseTest(unittest.TestCase):
+  """Tests TestCase functionality."""
+
+  def testTimeout(self):
+    """Test that test cases are interrupted when they are hanging."""
+
+    class TimeoutTestCase(cros_test_lib.TestCase):
+      """Test case that raises a TimeoutError because it takes too long."""
+
+      TEST_CASE_TIMEOUT = 1
+
+      def testSleeping(self):
+        """Sleep for 2 minutes. This should raise a TimeoutError."""
+        time.sleep(2 * 60)
+        raise AssertionError('Test case should have timed out.')
+
+    # Run the test case, verifying it raises a TimeoutError.
+    test = TimeoutTestCase(methodName='testSleeping')
+    self.assertRaises(cros_build_lib.TimeoutError, test.testSleeping)
+
+
 if __name__ == '__main__':
   cros_test_lib.main()