sdks/python/apache_beam/utils/retry_test.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""Unit tests for the retry module."""

# pytype: skip-file

import unittest

from parameterized import parameterized

from apache_beam.utils import retry

# Protect against environments where apitools library is not available.
# pylint: disable=wrong-import-order, wrong-import-position
# TODO(sourabhbajaj): Remove the GCP specific error code to a submodule
try:
  from apitools.base.py.exceptions import HttpError
except ImportError:
  HttpError = None
# pylint: enable=wrong-import-order, wrong-import-position


class FakeClock(object):
  """A fake clock object implementing sleep() and recording calls."""
  def __init__(self):
    self.calls = []

  def sleep(self, value):
    self.calls.append(value)


class FakeLogger(object):
  """A fake logger object implementing log() and recording calls."""
  def __init__(self):
    self.calls = []

  def log(self, message, interval, func_name, exn_name, exn_traceback):
    _ = interval, exn_traceback
    self.calls.append((message, func_name, exn_name))


@retry.with_exponential_backoff(clock=FakeClock())
def _test_function(a, b):
  _ = a, b
  raise NotImplementedError


@retry.with_exponential_backoff(initial_delay_secs=0.1, num_retries=1)
def _test_function_with_real_clock(a, b):
  _ = a, b
  raise NotImplementedError


@retry.no_retries
def _test_no_retry_function(a, b):
  _ = a, b
  raise NotImplementedError


class RetryTest(unittest.TestCase):
  def setUp(self):
    self.clock = FakeClock()
    self.logger = FakeLogger()
    self.calls = 0

  def permanent_failure(self, a, b):
    raise NotImplementedError

  def transient_failure(self, a, b):
    self.calls += 1
    if self.calls > 4:
      return a + b
    raise NotImplementedError

  def http_error(self, code):
    if HttpError is None:
      raise RuntimeError("This is not a valid test as GCP is not enabled")
    raise HttpError({'status': str(code)}, '', '')

  def test_with_explicit_decorator(self):
    # We pass one argument as positional argument and one as keyword argument
    # so that we cover both code paths for argument handling.
    self.assertRaises(NotImplementedError, _test_function, 10, b=20)

  def test_with_no_retry_decorator(self):
    self.assertRaises(NotImplementedError, _test_no_retry_function, 1, 2)

  def test_with_real_clock(self):
    self.assertRaises(
        NotImplementedError, _test_function_with_real_clock, 10, b=20)

  def test_with_default_number_of_retries(self):
    self.assertRaises(
        NotImplementedError,
        retry.with_exponential_backoff(clock=self.clock)(
            self.permanent_failure),
        10,
        b=20)
    self.assertEqual(len(self.clock.calls), 7)

  def test_with_explicit_number_of_retries(self):
    self.assertRaises(
        NotImplementedError,
        retry.with_exponential_backoff(clock=self.clock,
                                       num_retries=10)(self.permanent_failure),
        10,
        b=20)
    self.assertEqual(len(self.clock.calls), 10)

  @unittest.skipIf(HttpError is None, 'google-apitools is not installed')
  def test_with_http_error_that_should_not_be_retried(self):
    self.assertRaises(
        HttpError,
        retry.with_exponential_backoff(clock=self.clock,
                                       num_retries=10)(self.http_error),
        404)
    # Make sure just one call was made.
    self.assertEqual(len(self.clock.calls), 0)

  @unittest.skipIf(HttpError is None, 'google-apitools is not installed')
  def test_with_http_error_that_should_be_retried(self):
    self.assertRaises(
        HttpError,
        retry.with_exponential_backoff(clock=self.clock,
                                       num_retries=10)(self.http_error),
        500)
    self.assertEqual(len(self.clock.calls), 10)

  def test_with_explicit_initial_delay(self):
    self.assertRaises(
        NotImplementedError,
        retry.with_exponential_backoff(
            initial_delay_secs=10.0, clock=self.clock,
            fuzz=False)(self.permanent_failure),
        10,
        b=20)
    self.assertEqual(len(self.clock.calls), 7)
    self.assertEqual(self.clock.calls[0], 10.0)

  @parameterized.expand([(str(i), i) for i in range(0, 1000, 47)])
  def test_with_stop_after_secs(self, _, stop_after_secs):
    max_delay_secs = 10
    self.assertRaises(
        NotImplementedError,
        retry.with_exponential_backoff(
            num_retries=10000,
            initial_delay_secs=10.0,
            clock=self.clock,
            fuzz=False,
            max_delay_secs=max_delay_secs,
            stop_after_secs=stop_after_secs)(self.permanent_failure),
        10,
        b=20)
    total_delay = sum(self.clock.calls)
    self.assertLessEqual(total_delay, stop_after_secs)
    self.assertGreaterEqual(total_delay, stop_after_secs - max_delay_secs)

  def test_log_calls_for_permanent_failure(self):
    self.assertRaises(
        NotImplementedError,
        retry.with_exponential_backoff(
            clock=self.clock, logger=self.logger.log)(self.permanent_failure),
        10,
        b=20)
    self.assertEqual(len(self.logger.calls), 7)
    for message, func_name, exn_name in self.logger.calls:
      self.assertTrue(message.startswith('Retry with exponential backoff:'))
      self.assertEqual(exn_name, 'NotImplementedError\n')
      self.assertEqual(func_name, 'permanent_failure')

  def test_log_calls_for_transient_failure(self):
    result = retry.with_exponential_backoff(
        clock=self.clock, logger=self.logger.log,
        fuzz=False)(self.transient_failure)(
            10, b=20)
    self.assertEqual(result, 30)
    self.assertEqual(len(self.clock.calls), 4)
    self.assertEqual(self.clock.calls, [
        5.0 * 1,
        5.0 * 2,
        5.0 * 4,
        5.0 * 8,
    ])
    self.assertEqual(len(self.logger.calls), 4)
    for message, func_name, exn_name in self.logger.calls:
      self.assertTrue(message.startswith('Retry with exponential backoff:'))
      self.assertEqual(exn_name, 'NotImplementedError\n')
      self.assertEqual(func_name, 'transient_failure')


class DummyClass(object):
  def __init__(self, results):
    self.index = 0
    self.results = results

  @retry.with_exponential_backoff(num_retries=2, initial_delay_secs=0.1)
  def func(self):
    self.index += 1
    if self.index > len(self.results) or \
        self.results[self.index - 1] == "Error":
      raise ValueError("Error")
    return self.results[self.index - 1]


class RetryStateTest(unittest.TestCase):
  """The test_two_failures and test_single_failure would fail if we have
  any shared state for the retry decorator. This test tries to prevent a bug we
  found where the state in the decorator was shared across objects and retries
  were not available correctly.

  The test_call_two_objects would test this inside the same test.
  """
  def test_two_failures(self):
    dummy = DummyClass(["Error", "Error", "Success"])
    dummy.func()
    self.assertEqual(3, dummy.index)

  def test_single_failure(self):
    dummy = DummyClass(["Error", "Success"])
    dummy.func()
    self.assertEqual(2, dummy.index)

  def test_call_two_objects(self):
    dummy = DummyClass(["Error", "Error", "Success"])
    dummy.func()
    self.assertEqual(3, dummy.index)

    dummy2 = DummyClass(["Error", "Success"])
    dummy2.func()
    self.assertEqual(2, dummy2.index)


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