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