• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2015 The Chromium 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.
4import unittest
5
6import mock
7
8from py_utils import retry_util
9
10
11class RetryOnExceptionTest(unittest.TestCase):
12  def setUp(self):
13    self.num_calls = 0
14    # Patch time.sleep to make tests run faster (skip waits) and also check
15    # that exponential backoff is implemented correctly.
16    patcher = mock.patch('time.sleep')
17    self.time_sleep = patcher.start()
18    self.addCleanup(patcher.stop)
19
20  def testNoExceptionsReturnImmediately(self):
21    @retry_util.RetryOnException(Exception, retries=3)
22    def Test(retries=None):
23      del retries
24      self.num_calls += 1
25      return 'OK!'
26
27    # The function is called once and returns the expected value.
28    self.assertEqual(Test(), 'OK!')
29    self.assertEqual(self.num_calls, 1)
30
31  def testRaisesExceptionIfAlwaysFailing(self):
32    @retry_util.RetryOnException(KeyError, retries=5)
33    def Test(retries=None):
34      del retries
35      self.num_calls += 1
36      raise KeyError('oops!')
37
38    # The exception is eventually raised.
39    with self.assertRaises(KeyError):
40      Test()
41    # The function is called the expected number of times.
42    self.assertEqual(self.num_calls, 6)
43    # Waits between retries do follow exponential backoff.
44    self.assertEqual(
45        self.time_sleep.call_args_list,
46        [mock.call(i) for i in (1, 2, 4, 8, 16)])
47
48  def testOtherExceptionsAreNotCaught(self):
49    @retry_util.RetryOnException(KeyError, retries=3)
50    def Test(retries=None):
51      del retries
52      self.num_calls += 1
53      raise ValueError('oops!')
54
55    # The exception is raised immediately on the first try.
56    with self.assertRaises(ValueError):
57      Test()
58    self.assertEqual(self.num_calls, 1)
59
60  def testCallerMayOverrideRetries(self):
61    @retry_util.RetryOnException(KeyError, retries=3)
62    def Test(retries=None):
63      del retries
64      self.num_calls += 1
65      raise KeyError('oops!')
66
67    with self.assertRaises(KeyError):
68      Test(retries=10)
69    # The value on the caller overrides the default on the decorator.
70    self.assertEqual(self.num_calls, 11)
71
72  def testCanEventuallySucceed(self):
73    @retry_util.RetryOnException(KeyError, retries=5)
74    def Test(retries=None):
75      del retries
76      self.num_calls += 1
77      if self.num_calls < 3:
78        raise KeyError('oops!')
79      else:
80        return 'OK!'
81
82    # The value is returned after the expected number of calls.
83    self.assertEqual(Test(), 'OK!')
84    self.assertEqual(self.num_calls, 3)
85
86  def testRetriesCanBeSwitchedOff(self):
87    @retry_util.RetryOnException(KeyError, retries=5)
88    def Test(retries=None):
89      del retries
90      self.num_calls += 1
91      if self.num_calls < 3:
92        raise KeyError('oops!')
93      else:
94        return 'OK!'
95
96    # We fail immediately on the first try.
97    with self.assertRaises(KeyError):
98      Test(retries=0)
99    self.assertEqual(self.num_calls, 1)
100
101  def testCanRetryOnMultipleExceptions(self):
102    @retry_util.RetryOnException((KeyError, ValueError), retries=3)
103    def Test(retries=None):
104      del retries
105      self.num_calls += 1
106      if self.num_calls == 1:
107        raise KeyError('oops!')
108      elif self.num_calls == 2:
109        raise ValueError('uh oh!')
110      else:
111        return 'OK!'
112
113    # Call eventually succeeds after enough tries.
114    self.assertEqual(Test(retries=5), 'OK!')
115    self.assertEqual(self.num_calls, 3)
116
117
118if __name__ == '__main__':
119  unittest.main()
120