catapult/common/py_utils/py_utils/retry_util_unittest.py

# 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()