src/third_party/chromite/lib/timeout_util_unittest.py

   1 #!/usr/bin/python
   2 # Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
   3 # Use of this source code is governed by a BSD-style license that can be
   4 # found in the LICENSE file.
   5
   6 """Test suite for timeout_util.py"""
   7
   8 from __future__ import print_function
   9
  10 import datetime
  11 import os
  12 import sys
  13 import time
  14
  15 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(
  16     os.path.abspath(__file__)))))
  17
  18 from chromite.lib import cros_test_lib
  19 from chromite.lib import timeout_util
  20
  21
  22 # pylint: disable=W0212,R0904
  23
  24
  25 class TestTimeouts(cros_test_lib.TestCase):
  26   """Tests for timeout_util.Timeout."""
  27
  28   def testTimeout(self):
  29     """Tests that we can nest Timeout correctly."""
  30     self.assertFalse('mock' in str(time.sleep).lower())
  31     with timeout_util.Timeout(30):
  32       with timeout_util.Timeout(20):
  33         with timeout_util.Timeout(1):
  34           self.assertRaises(timeout_util.TimeoutError, time.sleep, 10)
  35
  36         # Should not raise a timeout exception as 20 > 2.
  37         time.sleep(1)
  38
  39   def testTimeoutNested(self):
  40     """Tests that we still re-raise an alarm if both are reached."""
  41     with timeout_util.Timeout(1):
  42       try:
  43         with timeout_util.Timeout(2):
  44           self.assertRaises(timeout_util.TimeoutError, time.sleep, 1)
  45
  46       # Craziness to catch nested timeouts.
  47       except timeout_util.TimeoutError:
  48         pass
  49       else:
  50         self.assertTrue(False, 'Should have thrown an exception')
  51
  52
  53 class TestWaitFors(cros_test_lib.TestCase):
  54   """Tests for assorted timeout_utils WaitForX methods."""
  55
  56   def setUp(self):
  57     self.values_ix = 0
  58     self.timestart = None
  59     self.timestop = None
  60
  61   def GetFunc(self, return_values):
  62     """Return a functor that returns given values in sequence with each call."""
  63     self.values_ix = 0
  64     self.timestart = None
  65     self.timestop = None
  66
  67     def _Func():
  68       if not self.timestart:
  69         self.timestart = datetime.datetime.utcnow()
  70
  71       val = return_values[self.values_ix]
  72       self.values_ix += 1
  73
  74       self.timestop = datetime.datetime.utcnow()
  75       return val
  76
  77     return _Func
  78
  79   def GetTryCount(self):
  80     """Get number of times func was tried."""
  81     return self.values_ix
  82
  83   def GetTrySeconds(self):
  84     """Get number of seconds that span all func tries."""
  85     delta = self.timestop - self.timestart
  86     return int(delta.seconds + 0.5)
  87
  88   def _TestWaitForSuccess(self, maxval, timeout, **kwargs):
  89     """Run through a test for WaitForSuccess."""
  90
  91     func = self.GetFunc(range(20))
  92     def _RetryCheck(val):
  93       return val < maxval
  94
  95     return timeout_util.WaitForSuccess(_RetryCheck, func, timeout, **kwargs)
  96
  97   def _TestWaitForReturnValue(self, values, timeout, **kwargs):
  98     """Run through a test for WaitForReturnValue."""
  99     func = self.GetFunc(range(20))
 100     return timeout_util.WaitForReturnValue(values, func, timeout, **kwargs)
 101
 102   def testWaitForSuccess1(self):
 103     """Test success after a few tries."""
 104     self.assertEquals(4, self._TestWaitForSuccess(4, 10, period=1))
 105     self.assertEquals(5, self.GetTryCount())
 106     self.assertEquals(4, self.GetTrySeconds())
 107
 108   def testWaitForSuccess2(self):
 109     """Test timeout after a couple tries."""
 110     self.assertRaises(timeout_util.TimeoutError, self._TestWaitForSuccess,
 111                       4, 3, period=1)
 112     self.assertEquals(3, self.GetTryCount())
 113     self.assertEquals(2, self.GetTrySeconds())
 114
 115   def testWaitForSuccess3(self):
 116     """Test success on first try."""
 117     self.assertEquals(0, self._TestWaitForSuccess(0, 10, period=1))
 118     self.assertEquals(1, self.GetTryCount())
 119     self.assertEquals(0, self.GetTrySeconds())
 120
 121   def testWaitForSuccess4(self):
 122     """Test success after a few tries with longer period."""
 123     self.assertEquals(3, self._TestWaitForSuccess(3, 10, period=2))
 124     self.assertEquals(4, self.GetTryCount())
 125     self.assertEquals(6, self.GetTrySeconds())
 126
 127   def testWaitForReturnValue1(self):
 128     """Test value found after a few tries."""
 129     self.assertEquals(4, self._TestWaitForReturnValue((4, 5), 10, period=1))
 130     self.assertEquals(5, self.GetTryCount())
 131     self.assertEquals(4, self.GetTrySeconds())
 132
 133   def testWaitForReturnValue2(self):
 134     """Test value found on first try."""
 135     self.assertEquals(0, self._TestWaitForReturnValue((0, 1), 10, period=1))
 136     self.assertEquals(1, self.GetTryCount())
 137     self.assertEquals(0, self.GetTrySeconds())
 138
 139
 140 if __name__ == '__main__':
 141   cros_test_lib.main()