aboutsummaryrefslogtreecommitdiff
path: root/catapult/common/py_utils/py_utils/retry_util_unittest.py
blob: 151f88ed8246452308878b2db9fba3f50354401e (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
# Copyright 2015 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
import mock
import unittest

from py_utils import retry_util


class RetryOnExceptionTest(unittest.TestCase):
  def setUp(self):
    self.num_calls = 0
    # Patch time.sleep to make tests run faster (skip waits) and also check
    # that exponential backoff is implemented correctly.
    patcher = mock.patch('time.sleep')
    self.time_sleep = patcher.start()
    self.addCleanup(patcher.stop)

  def testNoExceptionsReturnImmediately(self):
    @retry_util.RetryOnException(Exception, retries=3)
    def Test(retries=None):
      del retries
      self.num_calls += 1
      return 'OK!'

    # The function is called once and returns the expected value.
    self.assertEqual(Test(), 'OK!')
    self.assertEqual(self.num_calls, 1)

  def testRaisesExceptionIfAlwaysFailing(self):
    @retry_util.RetryOnException(KeyError, retries=5)
    def Test(retries=None):
      del retries
      self.num_calls += 1
      raise KeyError('oops!')

    # The exception is eventually raised.
    with self.assertRaises(KeyError):
      Test()
    # The function is called the expected number of times.
    self.assertEqual(self.num_calls, 6)
    # Waits between retries do follow exponential backoff.
    self.assertEqual(
        self.time_sleep.call_args_list,
        [mock.call(i) for i in (1, 2, 4, 8, 16)])

  def testOtherExceptionsAreNotCaught(self):
    @retry_util.RetryOnException(KeyError, retries=3)
    def Test(retries=None):
      del retries
      self.num_calls += 1
      raise ValueError('oops!')

    # The exception is raised immediately on the first try.
    with self.assertRaises(ValueError):
      Test()
    self.assertEqual(self.num_calls, 1)

  def testCallerMayOverrideRetries(self):
    @retry_util.RetryOnException(KeyError, retries=3)
    def Test(retries=None):
      del retries
      self.num_calls += 1
      raise KeyError('oops!')

    with self.assertRaises(KeyError):
      Test(retries=10)
    # The value on the caller overrides the default on the decorator.
    self.assertEqual(self.num_calls, 11)

  def testCanEventuallySucceed(self):
    @retry_util.RetryOnException(KeyError, retries=5)
    def Test(retries=None):
      del retries
      self.num_calls += 1
      if self.num_calls < 3:
        raise KeyError('oops!')
      else:
        return 'OK!'

    # The value is returned after the expected number of calls.
    self.assertEqual(Test(), 'OK!')
    self.assertEqual(self.num_calls, 3)

  def testRetriesCanBeSwitchedOff(self):
    @retry_util.RetryOnException(KeyError, retries=5)
    def Test(retries=None):
      del retries
      self.num_calls += 1
      if self.num_calls < 3:
        raise KeyError('oops!')
      else:
        return 'OK!'

    # We fail immediately on the first try.
    with self.assertRaises(KeyError):
      Test(retries=0)
    self.assertEqual(self.num_calls, 1)

  def testCanRetryOnMultipleExceptions(self):
    @retry_util.RetryOnException((KeyError, ValueError), retries=3)
    def Test(retries=None):
      del retries
      self.num_calls += 1
      if self.num_calls == 1:
        raise KeyError('oops!')
      elif self.num_calls == 2:
        raise ValueError('uh oh!')
      else:
        return 'OK!'

    # Call eventually succeeds after enough tries.
    self.assertEqual(Test(retries=5), 'OK!')
    self.assertEqual(self.num_calls, 3)


if __name__ == '__main__':
  unittest.main()