github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/utils/retry_test.py (about)

     1  #
     2  # Licensed to the Apache Software Foundation (ASF) under one or more
     3  # contributor license agreements.  See the NOTICE file distributed with
     4  # this work for additional information regarding copyright ownership.
     5  # The ASF licenses this file to You under the Apache License, Version 2.0
     6  # (the "License"); you may not use this file except in compliance with
     7  # the License.  You may obtain a copy of the License at
     8  #
     9  #    http://www.apache.org/licenses/LICENSE-2.0
    10  #
    11  # Unless required by applicable law or agreed to in writing, software
    12  # distributed under the License is distributed on an "AS IS" BASIS,
    13  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  # See the License for the specific language governing permissions and
    15  # limitations under the License.
    16  #
    17  
    18  """Unit tests for the retry module."""
    19  
    20  # pytype: skip-file
    21  
    22  import unittest
    23  
    24  from parameterized import parameterized
    25  
    26  from apache_beam.utils import retry
    27  
    28  # Protect against environments where apitools library is not available.
    29  # pylint: disable=wrong-import-order, wrong-import-position
    30  # TODO(sourabhbajaj): Remove the GCP specific error code to a submodule
    31  try:
    32    from apitools.base.py.exceptions import HttpError
    33  except ImportError:
    34    HttpError = None
    35  # pylint: enable=wrong-import-order, wrong-import-position
    36  
    37  
    38  class FakeClock(object):
    39    """A fake clock object implementing sleep() and recording calls."""
    40    def __init__(self):
    41      self.calls = []
    42  
    43    def sleep(self, value):
    44      self.calls.append(value)
    45  
    46  
    47  class FakeLogger(object):
    48    """A fake logger object implementing log() and recording calls."""
    49    def __init__(self):
    50      self.calls = []
    51  
    52    def log(self, message, interval, func_name, exn_name, exn_traceback):
    53      _ = interval, exn_traceback
    54      self.calls.append((message, func_name, exn_name))
    55  
    56  
    57  @retry.with_exponential_backoff(clock=FakeClock())
    58  def _test_function(a, b):
    59    _ = a, b
    60    raise NotImplementedError
    61  
    62  
    63  @retry.with_exponential_backoff(initial_delay_secs=0.1, num_retries=1)
    64  def _test_function_with_real_clock(a, b):
    65    _ = a, b
    66    raise NotImplementedError
    67  
    68  
    69  @retry.no_retries
    70  def _test_no_retry_function(a, b):
    71    _ = a, b
    72    raise NotImplementedError
    73  
    74  
    75  class RetryTest(unittest.TestCase):
    76    def setUp(self):
    77      self.clock = FakeClock()
    78      self.logger = FakeLogger()
    79      self.calls = 0
    80  
    81    def permanent_failure(self, a, b):
    82      raise NotImplementedError
    83  
    84    def transient_failure(self, a, b):
    85      self.calls += 1
    86      if self.calls > 4:
    87        return a + b
    88      raise NotImplementedError
    89  
    90    def http_error(self, code):
    91      if HttpError is None:
    92        raise RuntimeError("This is not a valid test as GCP is not enabled")
    93      raise HttpError({'status': str(code)}, '', '')
    94  
    95    def test_with_explicit_decorator(self):
    96      # We pass one argument as positional argument and one as keyword argument
    97      # so that we cover both code paths for argument handling.
    98      self.assertRaises(NotImplementedError, _test_function, 10, b=20)
    99  
   100    def test_with_no_retry_decorator(self):
   101      self.assertRaises(NotImplementedError, _test_no_retry_function, 1, 2)
   102  
   103    def test_with_real_clock(self):
   104      self.assertRaises(
   105          NotImplementedError, _test_function_with_real_clock, 10, b=20)
   106  
   107    def test_with_default_number_of_retries(self):
   108      self.assertRaises(
   109          NotImplementedError,
   110          retry.with_exponential_backoff(clock=self.clock)(
   111              self.permanent_failure),
   112          10,
   113          b=20)
   114      self.assertEqual(len(self.clock.calls), 7)
   115  
   116    def test_with_explicit_number_of_retries(self):
   117      self.assertRaises(
   118          NotImplementedError,
   119          retry.with_exponential_backoff(clock=self.clock,
   120                                         num_retries=10)(self.permanent_failure),
   121          10,
   122          b=20)
   123      self.assertEqual(len(self.clock.calls), 10)
   124  
   125    @unittest.skipIf(HttpError is None, 'google-apitools is not installed')
   126    def test_with_http_error_that_should_not_be_retried(self):
   127      self.assertRaises(
   128          HttpError,
   129          retry.with_exponential_backoff(clock=self.clock,
   130                                         num_retries=10)(self.http_error),
   131          404)
   132      # Make sure just one call was made.
   133      self.assertEqual(len(self.clock.calls), 0)
   134  
   135    @unittest.skipIf(HttpError is None, 'google-apitools is not installed')
   136    def test_with_http_error_that_should_be_retried(self):
   137      self.assertRaises(
   138          HttpError,
   139          retry.with_exponential_backoff(clock=self.clock,
   140                                         num_retries=10)(self.http_error),
   141          500)
   142      self.assertEqual(len(self.clock.calls), 10)
   143  
   144    def test_with_explicit_initial_delay(self):
   145      self.assertRaises(
   146          NotImplementedError,
   147          retry.with_exponential_backoff(
   148              initial_delay_secs=10.0, clock=self.clock,
   149              fuzz=False)(self.permanent_failure),
   150          10,
   151          b=20)
   152      self.assertEqual(len(self.clock.calls), 7)
   153      self.assertEqual(self.clock.calls[0], 10.0)
   154  
   155    @parameterized.expand([(str(i), i) for i in range(0, 1000, 47)])
   156    def test_with_stop_after_secs(self, _, stop_after_secs):
   157      max_delay_secs = 10
   158      self.assertRaises(
   159          NotImplementedError,
   160          retry.with_exponential_backoff(
   161              num_retries=10000,
   162              initial_delay_secs=10.0,
   163              clock=self.clock,
   164              fuzz=False,
   165              max_delay_secs=max_delay_secs,
   166              stop_after_secs=stop_after_secs)(self.permanent_failure),
   167          10,
   168          b=20)
   169      total_delay = sum(self.clock.calls)
   170      self.assertLessEqual(total_delay, stop_after_secs)
   171      self.assertGreaterEqual(total_delay, stop_after_secs - max_delay_secs)
   172  
   173    def test_log_calls_for_permanent_failure(self):
   174      self.assertRaises(
   175          NotImplementedError,
   176          retry.with_exponential_backoff(
   177              clock=self.clock, logger=self.logger.log)(self.permanent_failure),
   178          10,
   179          b=20)
   180      self.assertEqual(len(self.logger.calls), 7)
   181      for message, func_name, exn_name in self.logger.calls:
   182        self.assertTrue(message.startswith('Retry with exponential backoff:'))
   183        self.assertEqual(exn_name, 'NotImplementedError\n')
   184        self.assertEqual(func_name, 'permanent_failure')
   185  
   186    def test_log_calls_for_transient_failure(self):
   187      result = retry.with_exponential_backoff(
   188          clock=self.clock, logger=self.logger.log,
   189          fuzz=False)(self.transient_failure)(
   190              10, b=20)
   191      self.assertEqual(result, 30)
   192      self.assertEqual(len(self.clock.calls), 4)
   193      self.assertEqual(self.clock.calls, [
   194          5.0 * 1,
   195          5.0 * 2,
   196          5.0 * 4,
   197          5.0 * 8,
   198      ])
   199      self.assertEqual(len(self.logger.calls), 4)
   200      for message, func_name, exn_name in self.logger.calls:
   201        self.assertTrue(message.startswith('Retry with exponential backoff:'))
   202        self.assertEqual(exn_name, 'NotImplementedError\n')
   203        self.assertEqual(func_name, 'transient_failure')
   204  
   205  
   206  class DummyClass(object):
   207    def __init__(self, results):
   208      self.index = 0
   209      self.results = results
   210  
   211    @retry.with_exponential_backoff(num_retries=2, initial_delay_secs=0.1)
   212    def func(self):
   213      self.index += 1
   214      if self.index > len(self.results) or \
   215          self.results[self.index - 1] == "Error":
   216        raise ValueError("Error")
   217      return self.results[self.index - 1]
   218  
   219  
   220  class RetryStateTest(unittest.TestCase):
   221    """The test_two_failures and test_single_failure would fail if we have
   222    any shared state for the retry decorator. This test tries to prevent a bug we
   223    found where the state in the decorator was shared across objects and retries
   224    were not available correctly.
   225  
   226    The test_call_two_objects would test this inside the same test.
   227    """
   228    def test_two_failures(self):
   229      dummy = DummyClass(["Error", "Error", "Success"])
   230      dummy.func()
   231      self.assertEqual(3, dummy.index)
   232  
   233    def test_single_failure(self):
   234      dummy = DummyClass(["Error", "Success"])
   235      dummy.func()
   236      self.assertEqual(2, dummy.index)
   237  
   238    def test_call_two_objects(self):
   239      dummy = DummyClass(["Error", "Error", "Success"])
   240      dummy.func()
   241      self.assertEqual(3, dummy.index)
   242  
   243      dummy2 = DummyClass(["Error", "Success"])
   244      dummy2.func()
   245      self.assertEqual(2, dummy2.index)
   246  
   247  
   248  if __name__ == '__main__':
   249    unittest.main()