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
9import mox
10import time
11import unittest
12import signal
13
14import common
15from autotest_lib.client.common_lib.cros import retry
16from autotest_lib.client.common_lib import error
17
18
19class RetryTest(mox.MoxTestBase):
20    """Unit tests for retry decorators.
21
22    @var _FLAKY_FLAG: for use in tests that need to simulate random failures.
23    """
24
25    _FLAKY_FLAG = None
26
27    def setUp(self):
28        super(RetryTest, self).setUp()
29        self._FLAKY_FLAG = False
30
31
32    def testRetryDecoratorSucceeds(self):
33        """Tests that a wrapped function succeeds without retrying."""
34        @retry.retry(Exception)
35        def succeed():
36            return True
37
38        self.mox.StubOutWithMock(time, 'sleep')
39        self.mox.ReplayAll()
40        self.assertTrue(succeed())
41
42
43    def testRetryDecoratorFlakySucceeds(self):
44        """Tests that a wrapped function can retry and succeed."""
45        delay_sec = 10
46        @retry.retry(Exception, delay_sec=delay_sec)
47        def flaky_succeed():
48            if self._FLAKY_FLAG:
49                return True
50            self._FLAKY_FLAG = True
51            raise Exception()
52
53        self.mox.StubOutWithMock(time, 'sleep')
54        time.sleep(mox.Func(lambda x: abs(x - delay_sec) <= .5 * delay_sec))
55        self.mox.ReplayAll()
56        self.assertTrue(flaky_succeed())
57
58
59    def testRetryDecoratorFails(self):
60        """Tests that a wrapped function retries til the timeout, then fails."""
61        delay_sec = 10
62        @retry.retry(Exception, delay_sec=delay_sec)
63        def fail():
64            raise Exception()
65
66        self.mox.StubOutWithMock(time, 'sleep')
67        time.sleep(mox.Func(lambda x: abs(x - delay_sec) <= .5 * delay_sec))
68        self.mox.ReplayAll()
69        self.assertRaises(Exception, fail)
70
71
72    def testRetryDecoratorRaisesCrosDynamicSuiteException(self):
73        """Tests that dynamic_suite exceptions raise immediately, no retry."""
74        @retry.retry(Exception)
75        def fail():
76            raise error.ControlFileNotFound()
77
78        self.mox.StubOutWithMock(time, 'sleep')
79        self.mox.ReplayAll()
80        self.assertRaises(error.ControlFileNotFound, fail)
81
82
83    def testRetryDecoratorFailsWithTimeout(self):
84        """Tests that a wrapped function retries til the timeout, then fails."""
85        @retry.retry(Exception, timeout_min=0.02, delay_sec=0.1)
86        def fail():
87            time.sleep(2)
88            return True
89
90        self.mox.ReplayAll()
91        #self.assertEquals(None, fail())
92        self.assertRaises(error.TimeoutException, fail)
93
94
95    def testRetryDecoratorSucceedsBeforeTimeout(self):
96        """Tests that a wrapped function succeeds before the timeout."""
97        @retry.retry(Exception, timeout_min=0.02, delay_sec=0.1)
98        def succeed():
99            time.sleep(0.1)
100            return True
101
102        self.mox.ReplayAll()
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.mox.ReplayAll()
131        self.assertRaises(TestTimeoutException, testFunc)
132
133
134    def testRetryDecoratorWithNoAlarmLeak(self):
135        """Tests that a wrapped function throws exception before the timeout
136        and no signal is leaked."""
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.mox.ReplayAll()
163        self.assertFalse(testFunc())
164
165
166if __name__ == '__main__':
167    unittest.main()
168