#
# 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."""

from __future__ import absolute_import

import unittest
from builtins import object

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)

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