Home | History | Annotate | Download | only in cros 
      1 #!/usr/bin/env python
      2 
      3 # Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
      4 # Use of this source code is governed by a BSD-style license that can be
      5 # found in the LICENSE file.
      6 
      7 """Unit tests for client/common_lib/cros/retry.py."""
      8 
      9 import itertools
     10 import mox
     11 import time
     12 import unittest
     13 import signal
     14 
     15 import mock
     16 
     17 import common
     18 from autotest_lib.client.common_lib.cros import retry
     19 from autotest_lib.client.common_lib import error
     20 
     21 
     22 class RetryTest(mox.MoxTestBase):
     23     """Unit tests for retry decorators.
     24 
     25     @var _FLAKY_FLAG: for use in tests that need to simulate random failures.
     26     """
     27 
     28     _FLAKY_FLAG = None
     29 
     30     def setUp(self):
     31         super(RetryTest, self).setUp()
     32         self._FLAKY_FLAG = False
     33 
     34         patcher = mock.patch('time.sleep', autospec=True)
     35         self._sleep_mock = patcher.start()
     36         self.addCleanup(patcher.stop)
     37 
     38         patcher = mock.patch('time.time', autospec=True)
     39         self._time_mock = patcher.start()
     40         self.addCleanup(patcher.stop)
     41 
     42 
     43     def testRetryDecoratorSucceeds(self):
     44         """Tests that a wrapped function succeeds without retrying."""
     45         @retry.retry(Exception)
     46         def succeed():
     47             return True
     48         self.assertTrue(succeed())
     49         self.assertFalse(self._sleep_mock.called)
     50 
     51 
     52     def testRetryDecoratorFlakySucceeds(self):
     53         """Tests that a wrapped function can retry and succeed."""
     54         delay_sec = 10
     55         self._time_mock.side_effect = itertools.count(delay_sec)
     56         @retry.retry(Exception, delay_sec=delay_sec)
     57         def flaky_succeed():
     58             if self._FLAKY_FLAG:
     59                 return True
     60             self._FLAKY_FLAG = True
     61             raise Exception()
     62         self.assertTrue(flaky_succeed())
     63 
     64 
     65     def testRetryDecoratorFails(self):
     66         """Tests that a wrapped function retries til the timeout, then fails."""
     67         delay_sec = 10
     68         self._time_mock.side_effect = itertools.count(delay_sec)
     69         @retry.retry(Exception, delay_sec=delay_sec)
     70         def fail():
     71             raise Exception()
     72         self.assertRaises(Exception, fail)
     73 
     74 
     75     def testRetryDecoratorRaisesCrosDynamicSuiteException(self):
     76         """Tests that dynamic_suite exceptions raise immediately, no retry."""
     77         @retry.retry(Exception)
     78         def fail():
     79             raise error.ControlFileNotFound()
     80         self.assertRaises(error.ControlFileNotFound, fail)
     81 
     82 
     83 
     84 
     85 class ActualRetryTest(unittest.TestCase):
     86     """Unit tests for retry decorators with real sleep."""
     87 
     88     def testRetryDecoratorFailsWithTimeout(self):
     89         """Tests that a wrapped function retries til the timeout, then fails."""
     90         @retry.retry(Exception, timeout_min=0.02, delay_sec=0.1)
     91         def fail():
     92             time.sleep(2)
     93             return True
     94         self.assertRaises(error.TimeoutException, fail)
     95 
     96 
     97     def testRetryDecoratorSucceedsBeforeTimeout(self):
     98         """Tests that a wrapped function succeeds before the timeout."""
     99         @retry.retry(Exception, timeout_min=0.02, delay_sec=0.1)
    100         def succeed():
    101             time.sleep(0.1)
    102             return True
    103         self.assertTrue(succeed())
    104 
    105 
    106     def testRetryDecoratorSucceedsWithExistingSignal(self):
    107         """Tests that a wrapped function succeeds before the timeout and
    108         previous signal being restored."""
    109         class TestTimeoutException(Exception):
    110             pass
    111 
    112         def testFunc():
    113             @retry.retry(Exception, timeout_min=0.05, delay_sec=0.1)
    114             def succeed():
    115                 time.sleep(0.1)
    116                 return True
    117 
    118             succeed()
    119             # Wait for 1.5 second for previous signal to be raised
    120             time.sleep(1.5)
    121 
    122         def testHandler(signum, frame):
    123             """
    124             Register a handler for the timeout.
    125             """
    126             raise TestTimeoutException('Expected timed out.')
    127 
    128         signal.signal(signal.SIGALRM, testHandler)
    129         signal.alarm(1)
    130         self.assertRaises(TestTimeoutException, testFunc)
    131 
    132 
    133     def testRetryDecoratorWithNoAlarmLeak(self):
    134         """Tests that a wrapped function throws exception before the timeout
    135         and no signal is leaked."""
    136 
    137         def testFunc():
    138             @retry.retry(Exception, timeout_min=0.06, delay_sec=0.1)
    139             def fail():
    140                 time.sleep(0.1)
    141                 raise Exception()
    142 
    143 
    144             def testHandler(signum, frame):
    145                 """
    146                 Register a handler for the timeout.
    147                 """
    148                 self.alarm_leaked = True
    149 
    150 
    151             # Set handler for signal.SIGALRM to catch any leaked alarm.
    152             self.alarm_leaked = False
    153             signal.signal(signal.SIGALRM, testHandler)
    154             try:
    155                 fail()
    156             except Exception:
    157                 pass
    158             # Wait for 2 seconds to check if any alarm is leaked
    159             time.sleep(2)
    160             return self.alarm_leaked
    161 
    162         self.assertFalse(testFunc())
    163 
    164 
    165 if __name__ == '__main__':
    166     unittest.main()
    167