/*
 * 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.
 */
package org.apache.beam.sdk.extensions.gcp.util;

import static org.hamcrest.MatcherAssert.assertThat;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.fail;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Matchers.any;
import static org.mockito.Matchers.anyInt;
import static org.mockito.Mockito.atLeastOnce;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoMoreInteractions;
import static org.mockito.Mockito.when;

import com.google.api.client.http.HttpResponse;
import com.google.api.client.http.HttpResponseException;
import com.google.api.client.http.HttpResponseInterceptor;
import com.google.api.client.http.HttpTransport;
import com.google.api.client.http.LowLevelHttpRequest;
import com.google.api.client.http.LowLevelHttpResponse;
import com.google.api.client.json.JsonFactory;
import com.google.api.client.json.jackson2.JacksonFactory;
import com.google.api.client.testing.http.MockHttpTransport;
import com.google.api.client.testing.http.MockLowLevelHttpRequest;
import com.google.api.client.testing.http.MockLowLevelHttpResponse;
import com.google.api.client.util.NanoClock;
import com.google.api.services.storage.Storage;
import com.google.api.services.storage.Storage.Objects.Get;
import java.io.IOException;
import java.net.SocketTimeoutException;
import java.security.PrivateKey;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.atomic.AtomicLong;
import org.apache.beam.sdk.testing.ExpectedLogs;
import org.apache.beam.sdk.util.FastNanoClockAndSleeper;
import org.hamcrest.Matchers;
import org.junit.After;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;

/** Tests for RetryHttpRequestInitializer. */
@RunWith(JUnit4.class)
public class RetryHttpRequestInitializerTest {

  @Rule public ExpectedLogs expectedLogs = ExpectedLogs.none(RetryHttpRequestInitializer.class);

  @Mock private PrivateKey mockPrivateKey;
  @Mock private LowLevelHttpRequest mockLowLevelRequest;
  @Mock private LowLevelHttpResponse mockLowLevelResponse;
  @Mock private HttpResponseInterceptor mockHttpResponseInterceptor;

  private final JsonFactory jsonFactory = JacksonFactory.getDefaultInstance();
  private Storage storage;

  // Used to test retrying a request more than the default 10 times.
  static class MockNanoClock implements NanoClock {
    private int[] timesMs = {
      500, 750, 1125, 1688, 2531, 3797, 5695, 8543, 12814, 19222, 28833, 43249, 64873, 97310,
      145965, 218945, 328420
    };
    private int i = 0;

    @Override
    public long nanoTime() {
      return timesMs[i++ / 2] * 1000000L;
    }
  }

  MockLowLevelHttpResponse[] createMockResponseWithStatusCode(int... statusCodes) {
    MockLowLevelHttpResponse[] responses = new MockLowLevelHttpResponse[statusCodes.length];

    for (int i = 0; i < statusCodes.length; ++i) {
      MockLowLevelHttpResponse response = mock(MockLowLevelHttpResponse.class);
      when(response.getStatusCode()).thenReturn(statusCodes[i]);
      responses[i] = response;
    }
    return responses;
  }

  @Before
  public void setUp() {
    MockitoAnnotations.initMocks(this);

    HttpTransport lowLevelTransport =
        new HttpTransport() {
          @Override
          protected LowLevelHttpRequest buildRequest(String method, String url) throws IOException {
            return mockLowLevelRequest;
          }
        };

    // Retry initializer will pass through to credential, since we can have
    // only a single HttpRequestInitializer, and we use multiple Credential
    // types in the SDK, not all of which allow for retry configuration.
    RetryHttpRequestInitializer initializer =
        new RetryHttpRequestInitializer(
            new MockNanoClock(),
            millis -> {},
            Arrays.asList(418 /* I'm a teapot */),
            mockHttpResponseInterceptor);
    storage =
        new Storage.Builder(lowLevelTransport, jsonFactory, initializer)
            .setApplicationName("test")
            .build();
  }

  @After
  public void tearDown() {
    verifyNoMoreInteractions(mockPrivateKey);
    verifyNoMoreInteractions(mockLowLevelRequest);
    verifyNoMoreInteractions(mockHttpResponseInterceptor);
  }

  @Test
  public void testBasicOperation() throws IOException {
    when(mockLowLevelRequest.execute()).thenReturn(mockLowLevelResponse);
    when(mockLowLevelResponse.getStatusCode()).thenReturn(200);

    Storage.Buckets.Get result = storage.buckets().get("test");
    HttpResponse response = result.executeUnparsed();
    assertNotNull(response);

    verify(mockHttpResponseInterceptor).interceptResponse(any(HttpResponse.class));
    verify(mockLowLevelRequest, atLeastOnce()).addHeader(anyString(), anyString());
    verify(mockLowLevelRequest).setTimeout(anyInt(), anyInt());
    verify(mockLowLevelRequest).setWriteTimeout(anyInt());
    verify(mockLowLevelRequest).execute();
    verify(mockLowLevelResponse, atLeastOnce()).getStatusCode();
    expectedLogs.verifyNotLogged("Request failed");
  }

  /** Tests that a non-retriable error is not retried. */
  @Test
  public void testErrorCodeForbidden() throws IOException {
    MockLowLevelHttpResponse[] responses =
        createMockResponseWithStatusCode(
            403, // Non-retryable error.
            200); // Shouldn't happen.

    when(mockLowLevelRequest.execute()).thenReturn(responses[0], responses[1]);

    try {
      Storage.Buckets.Get result = storage.buckets().get("test");
      HttpResponse response = result.executeUnparsed();
      assertNotNull(response);
    } catch (HttpResponseException e) {
      assertThat(e.getMessage(), Matchers.containsString("403"));
    }

    verify(mockHttpResponseInterceptor).interceptResponse(any(HttpResponse.class));
    verify(mockLowLevelRequest, atLeastOnce()).addHeader(anyString(), anyString());
    verify(mockLowLevelRequest).setTimeout(anyInt(), anyInt());
    verify(mockLowLevelRequest).setWriteTimeout(anyInt());
    verify(mockLowLevelRequest).execute();
    verify(responses[0], atLeastOnce()).getStatusCode();
    verify(responses[1], never()).getStatusCode();
    expectedLogs.verifyWarn("Request failed with code 403");
  }

  /** Tests that a retriable error is retried. */
  @Test
  public void testRetryableError() throws IOException {
    MockLowLevelHttpResponse[] mockResponses =
        createMockResponseWithStatusCode(
            503, // Retryable
            429, // We also retry on 429 Too Many Requests.
            200);
    when(mockLowLevelRequest.execute())
        .thenReturn(mockResponses[0], mockResponses[1], mockResponses[2]);

    Storage.Buckets.Get result = storage.buckets().get("test");
    HttpResponse response = result.executeUnparsed();
    assertNotNull(response);

    verify(mockHttpResponseInterceptor).interceptResponse(any(HttpResponse.class));
    verify(mockLowLevelRequest, atLeastOnce()).addHeader(anyString(), anyString());
    verify(mockLowLevelRequest, times(3)).setTimeout(anyInt(), anyInt());
    verify(mockLowLevelRequest, times(3)).setWriteTimeout(anyInt());
    verify(mockLowLevelRequest, times(3)).execute();

    // It reads the status code of all responses
    for (MockLowLevelHttpResponse mockResponse : mockResponses) {
      verify(mockResponse, atLeastOnce()).getStatusCode();
    }
    expectedLogs.verifyDebug("Request failed with code 503");
  }

  /** Tests that an IOException is retried. */
  @Test
  public void testThrowIOException() throws IOException {
    when(mockLowLevelRequest.execute())
        .thenThrow(new IOException("Fake Error"))
        .thenReturn(mockLowLevelResponse);
    when(mockLowLevelResponse.getStatusCode()).thenReturn(200);

    Storage.Buckets.Get result = storage.buckets().get("test");
    HttpResponse response = result.executeUnparsed();
    assertNotNull(response);

    verify(mockHttpResponseInterceptor).interceptResponse(any(HttpResponse.class));
    verify(mockLowLevelRequest, atLeastOnce()).addHeader(anyString(), anyString());
    verify(mockLowLevelRequest, times(2)).setTimeout(anyInt(), anyInt());
    verify(mockLowLevelRequest, times(2)).setWriteTimeout(anyInt());
    verify(mockLowLevelRequest, times(2)).execute();
    verify(mockLowLevelResponse, atLeastOnce()).getStatusCode();
    expectedLogs.verifyDebug("Request failed with IOException");
  }

  /** Tests that a retryable error is retried enough times. */
  @Test
  public void testRetryableErrorRetryEnoughTimes() throws IOException {
    List<MockLowLevelHttpResponse> responses = new ArrayList<>();
    final int retries = 10;

    // The underlying http library calls getStatusCode method of a response multiple times. For a
    // response, the method should return the same value. Therefore this test cannot rely on
    // `mockLowLevelResponse` variable that are reused across responses.
    when(mockLowLevelRequest.execute())
        .thenAnswer(
            new Answer<MockLowLevelHttpResponse>() {
              int n = 0;

              @Override
              public MockLowLevelHttpResponse answer(InvocationOnMock invocation) throws Throwable {
                MockLowLevelHttpResponse response = mock(MockLowLevelHttpResponse.class);
                responses.add(response);
                when(response.getStatusCode()).thenReturn(n++ < retries ? 503 : 9999);
                return response;
              }
            });

    Storage.Buckets.Get result = storage.buckets().get("test");
    try {
      result.executeUnparsed();
      fail();
    } catch (IOException e) {
    }

    verify(mockHttpResponseInterceptor).interceptResponse(any(HttpResponse.class));
    verify(mockLowLevelRequest, atLeastOnce()).addHeader(anyString(), anyString());
    verify(mockLowLevelRequest, times(retries + 1)).setTimeout(anyInt(), anyInt());
    verify(mockLowLevelRequest, times(retries + 1)).setWriteTimeout(anyInt());
    verify(mockLowLevelRequest, times(retries + 1)).execute();
    assertThat(responses, Matchers.hasSize(retries + 1));
    for (MockLowLevelHttpResponse response : responses) {
      verify(response, atLeastOnce()).getStatusCode();
    }
    expectedLogs.verifyWarn("performed 10 retries due to unsuccessful status codes");
  }

  /**
   * Tests that when RPCs fail with {@link SocketTimeoutException}, the IO exception handler is
   * invoked.
   */
  @Test
  @SuppressWarnings("AssertionFailureIgnored")
  public void testIOExceptionHandlerIsInvokedOnTimeout() throws Exception {
    FastNanoClockAndSleeper fakeClockAndSleeper = new FastNanoClockAndSleeper();
    // Counts the number of calls to execute the HTTP request.
    final AtomicLong executeCount = new AtomicLong();

    // 10 is a private internal constant in the Google API Client library. See
    // com.google.api.client.http.HttpRequest#setNumberOfRetries
    // TODO: update this test once the private internal constant is public.
    final int defaultNumberOfRetries = 10;

    // A mock HTTP request that always throws SocketTimeoutException.
    MockHttpTransport transport =
        new MockHttpTransport.Builder()
            .setLowLevelHttpRequest(
                new MockLowLevelHttpRequest() {
                  @Override
                  public LowLevelHttpResponse execute() throws IOException {
                    executeCount.incrementAndGet();
                    throw new SocketTimeoutException("Fake forced timeout exception");
                  }
                })
            .build();

    // A sample HTTP request to Google Cloud Storage that uses both a default Transport and
    // effectively a default RetryHttpRequestInitializer (same args as default with fake
    // clock/sleeper).
    Storage storage =
        new Storage.Builder(
                transport,
                Transport.getJsonFactory(),
                new RetryHttpRequestInitializer(
                    fakeClockAndSleeper::nanoTime,
                    fakeClockAndSleeper::sleep,
                    Collections.emptyList(),
                    null))
            .build();

    Get getRequest = storage.objects().get("gs://fake", "file");

    Throwable thrown = null;
    try {
      getRequest.execute();
    } catch (Throwable e) {
      thrown = e;
    }
    assertNotNull("Expected execute to throw an exception", thrown);
    assertThat(thrown, Matchers.instanceOf(SocketTimeoutException.class));
    assertEquals(1 + defaultNumberOfRetries, executeCount.get());
    expectedLogs.verifyWarn("performed 10 retries due to IOExceptions");
  }
}
