com.google.api.client.testing.util.MockSleeper Java Examples

@Test
public void testHttpRetries() throws Exception {
  List<MockLowLevelHttpResponse> mockResponses = ImmutableList.of(
      firstRpcResponse.setStatusCode(503).setContent("{}"),
      new MockLowLevelHttpResponse().setContent("{}"));
  MockHttpTransport transport = new MultiRequestMockHttpTransport(mockResponses);
  FirebaseOptions options = new FirebaseOptions.Builder()
      .setCredentials(new MockGoogleCredentials("test-token"))
      .setProjectId(PROJECT_ID)
      .setHttpTransport(transport)
      .build();
  FirebaseApp app = FirebaseApp.initializeApp(options);
  HttpRequestFactory requestFactory = TestApiClientUtils.delayBypassedRequestFactory(app);
  FirebaseProjectManagementServiceImpl serviceImpl = new FirebaseProjectManagementServiceImpl(
      app, new MockSleeper(), new MockScheduler(), requestFactory);
  serviceImpl.setInterceptor(interceptor);

  serviceImpl.deleteShaCertificate(SHA1_RESOURCE_NAME);

  String expectedUrl = String.format(
      "%s/v1beta1/%s", FIREBASE_PROJECT_MANAGEMENT_URL, SHA1_RESOURCE_NAME);
  checkRequestHeader(expectedUrl, HttpMethod.DELETE);
}
 

@Test
public void testCallRetriesAndSucceeds() throws Exception {
  MockSleeper sleeper = new MockSleeper();
  ArrayList<Exception> exceptions = new ArrayList<>();
  exceptions.add(new SocketTimeoutException("socket"));
  exceptions.add(new SocketTimeoutException("socket2"));
  exceptions.add(new SocketTimeoutException("socket3"));
  CallableTester<Exception> callTester = new CallableTester<>(exceptions);
  BackOff backoff = new RetryBoundedBackOff(3, new BackOffTester());
  assertThat(
          ResilientOperation.retry(
              callTester, backoff, RetryDeterminer.DEFAULT, Exception.class, sleeper))
      .isEqualTo(3);
  assertThat(callTester.timesCalled()).isEqualTo(4);
  verifySleeper(sleeper, 3);
}
 

@Test
public void testCallMaxRetries() throws Exception {
  MockSleeper sleeper = new MockSleeper();
  ArrayList<Exception> exceptions = new ArrayList<>();
  exceptions.add(new SocketTimeoutException("socket"));
  exceptions.add(new SocketTimeoutException("socket2"));
  exceptions.add(new SocketTimeoutException("socket3"));
  CallableTester<Exception> callTester = new CallableTester<>(exceptions);
  BackOff backoff = new RetryBoundedBackOff(2, new BackOffTester());

  SocketTimeoutException thrown =
      assertThrows(
          SocketTimeoutException.class,
          () ->
              ResilientOperation.retry(
                  callTester, backoff, RetryDeterminer.DEFAULT, Exception.class, sleeper));
  assertThat(thrown).hasMessageThat().contains("socket3");

  assertThat(callTester.timesCalled()).isEqualTo(3);
  verifySleeper(sleeper, 2);
}
 

@Test
public void testRetryOnIOException() throws IOException {
  MockSleeper sleeper = new MockSleeper();
  RetryInitializer initializer = new RetryInitializer(
      retryOnIOAndServiceUnavailableErrors(sleeper));

  CountingLowLevelHttpRequest failingRequest = CountingLowLevelHttpRequest.fromException(
      new IOException("test error"));
  HttpRequest request = TestUtils.createRequest(failingRequest);
  initializer.initialize(request);
  final HttpUnsuccessfulResponseHandler retryHandler = request.getUnsuccessfulResponseHandler();

  try {
    request.execute();
    fail("No exception thrown for HTTP error");
  } catch (IOException e) {
    assertEquals("test error", e.getMessage());
  }

  assertEquals(MAX_RETRIES, sleeper.getCount());
  assertEquals(MAX_RETRIES + 1, failingRequest.getCount());
  assertSame(retryHandler, request.getUnsuccessfulResponseHandler());
}
 

@Test
public void testRetryOnHttpError() throws IOException {
  MockSleeper sleeper = new MockSleeper();
  RetryInitializer initializer = new RetryInitializer(
      retryOnIOAndServiceUnavailableErrors(sleeper));
  CountingLowLevelHttpRequest failingRequest = CountingLowLevelHttpRequest.fromStatus(503);
  HttpRequest request = TestUtils.createRequest(failingRequest);
  initializer.initialize(request);
  final HttpUnsuccessfulResponseHandler retryHandler = request.getUnsuccessfulResponseHandler();

  try {
    request.execute();
    fail("No exception thrown for HTTP error");
  } catch (HttpResponseException e) {
    assertEquals(503, e.getStatusCode());
  }

  assertEquals(MAX_RETRIES, sleeper.getCount());
  assertEquals(MAX_RETRIES + 1, failingRequest.getCount());
  assertSame(retryHandler, request.getUnsuccessfulResponseHandler());
}
 

@Test
public void testDoesNotRetryAfterInterruption() throws IOException {
  MockSleeper sleeper = new MockSleeper() {
    @Override
    public void sleep(long millis) throws InterruptedException {
      super.sleep(millis);
      throw new InterruptedException();
    }
  };
  RetryUnsuccessfulResponseHandler handler = new RetryUnsuccessfulResponseHandler(
      testRetryConfig(sleeper));
  CountingLowLevelHttpRequest failingRequest = CountingLowLevelHttpRequest.fromStatus(503);
  HttpRequest request = TestUtils.createRequest(failingRequest);
  request.setUnsuccessfulResponseHandler(handler);
  request.setNumberOfRetries(MAX_RETRIES);

  try {
    request.execute();
    fail("No exception thrown for HTTP error");
  } catch (HttpResponseException e) {
    assertEquals(503, e.getStatusCode());
  }

  assertEquals(1, sleeper.getCount());
  assertEquals(1, failingRequest.getCount());
}
 

@Test
public void testCallRetriesAndFailsWithSocketErrors() throws Exception {
  MockSleeper sleeper = new MockSleeper();
  ArrayList<IOException> exceptions = new ArrayList<>();
  exceptions.add(new SocketTimeoutException("socket"));
  exceptions.add(new SocketTimeoutException("socket"));
  exceptions.add(new IOException("FakeException"));
  CallableTester<IOException> callTester = new CallableTester<>(exceptions);
  BackOff backoff = new RetryBoundedBackOff(5, new BackOffTester());

  IOException thrown =
      assertThrows(
          IOException.class,
          () ->
              ResilientOperation.retry(
                  callTester,
                  backoff,
                  RetryDeterminer.SOCKET_ERRORS,
                  IOException.class,
                  sleeper));
  assertThat(thrown).hasMessageThat().contains("FakeException");

  assertThat(callTester.timesCalled()).isEqualTo(3);
  verifySleeper(sleeper, 2);
}
 

@Test
public void testCallRetriesAndFails() throws Exception {
  MockSleeper sleeper = new MockSleeper();
  ArrayList<Exception> exceptions = new ArrayList<>();
  exceptions.add(new SocketTimeoutException("socket"));
  exceptions.add(new SocketTimeoutException("socket"));
  exceptions.add(new IllegalArgumentException("FakeException"));
  CallableTester<Exception> callTester = new CallableTester<>(exceptions);
  BackOff backoff = new RetryBoundedBackOff(5, new BackOffTester());

  IllegalArgumentException thrown =
      assertThrows(
          IllegalArgumentException.class,
          () ->
              ResilientOperation.retry(
                  callTester, backoff, RetryDeterminer.DEFAULT, Exception.class, sleeper));
  assertThat(thrown).hasMessageThat().contains("FakeException");

  assertThat(callTester.timesCalled()).isEqualTo(3);
  verifySleeper(sleeper, 2);
}
 

@Test
public void testCallFailsOnBadException() throws Exception {
  MockSleeper sleeper = new MockSleeper();
  ArrayList<Exception> exceptions = new ArrayList<>();
  exceptions.add(new IllegalArgumentException("FakeException"));
  CallableTester<Exception> callTester = new CallableTester<>(exceptions);
  BackOff backoff = new RetryBoundedBackOff(3, new BackOffTester());

  IllegalArgumentException thrown =
      assertThrows(
          IllegalArgumentException.class,
          () ->
              ResilientOperation.retry(
                  callTester, backoff, RetryDeterminer.DEFAULT, Exception.class, sleeper));
  assertThat(thrown).hasMessageThat().contains("FakeException");

  assertThat(callTester.timesCalled()).isEqualTo(1);
  verifySleeper(sleeper, 0);
}
 

public void subtestNext(long millis) throws Exception {
  MockSleeper sleeper = new MockSleeper();
  MockBackOff backOff = new MockBackOff().setBackOffMillis(millis);
  if (millis == BackOff.STOP) {
    backOff.setMaxTries(0);
  }
  int retries = 0;
  while (retries <= backOff.getMaxTries() + 1) {
    boolean next = BackOffUtils.next(sleeper, backOff);
    assertEquals(retries < backOff.getMaxTries(), next);
    if (next) {
      retries++;
    }
    assertEquals(retries, sleeper.getCount());
    assertEquals(retries, backOff.getNumberOfTries());
    if (!next) {
      break;
    }
    assertEquals(millis, sleeper.getLastMillis());
  }
}
 

private static FirebaseProjectManagementServiceImpl initServiceImpl(
    List<MockLowLevelHttpResponse> mockResponses,
    MultiRequestTestResponseInterceptor interceptor) {
  MockHttpTransport transport = new MultiRequestMockHttpTransport(mockResponses);
  FirebaseOptions options = new FirebaseOptions.Builder()
      .setCredentials(new MockGoogleCredentials("test-token"))
      .setProjectId(PROJECT_ID)
      .setHttpTransport(transport)
      .build();
  FirebaseApp app = FirebaseApp.initializeApp(options);
  HttpRequestFactory requestFactory = TestApiClientUtils.retryDisabledRequestFactory(app);
  FirebaseProjectManagementServiceImpl serviceImpl = new FirebaseProjectManagementServiceImpl(
      app, new MockSleeper(), new MockScheduler(), requestFactory);
  serviceImpl.setInterceptor(interceptor);
  return serviceImpl;
}
 

public void verifySleeper(MockSleeper sleeper, int retry) {
  assertThat(retry).isEqualTo(sleeper.getCount());
  if (retry == 0) {
    return;
  }
  assertThat((long) Math.pow(2, retry)).isEqualTo(sleeper.getLastMillis());
}
 

@Test
public void testValidCallHasNoRetries() throws Exception {
  MockSleeper sleeper = new MockSleeper();
  CallableTester<Exception> callTester = new CallableTester<>(new ArrayList<Exception>());
  BackOff backoff = new RetryBoundedBackOff(3, new BackOffTester());
  ResilientOperation.retry(callTester, backoff, RetryDeterminer.DEFAULT, Exception.class,
      sleeper);
  assertThat(callTester.timesCalled()).isEqualTo(1);
  assertThat(sleeper.getCount()).isEqualTo(0);
}
 

private void subsetHandleResponse(
    int count, int millis, boolean retry, BackOff backOff, BackOffRequired backOffRequired)
    throws IOException {
  // create the handler
  MockSleeper sleeper = new MockSleeper();
  HttpBackOffUnsuccessfulResponseHandler handler =
      new HttpBackOffUnsuccessfulResponseHandler(backOff)
          .setSleeper(sleeper)
          .setBackOffRequired(backOffRequired);

  while (handler.handleResponse(null, null, retry)) {
    assertEquals(millis, sleeper.getLastMillis());
  }
  assertEquals(count, sleeper.getCount());
}
 

public void subsetHandle(long count, long millis, boolean retrySupported, BackOff backOff)
    throws IOException {
  // create the handler
  MockSleeper sleeper = new MockSleeper();
  HttpBackOffIOExceptionHandler handler =
      new HttpBackOffIOExceptionHandler(backOff).setSleeper(sleeper);

  while (handler.handleIOException(null, retrySupported)) {
    assertEquals(millis, sleeper.getLastMillis());
  }
  assertEquals(count, sleeper.getCount());
}
 

/**
 * Tests that {@link DatasetServiceImpl#insertAll} retries other non-rate-limited,
 * non-quota-exceeded attempts.
 */
@Test
public void testInsertOtherRetry() throws Throwable {
  TableReference ref =
      new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
  List<ValueInSingleWindow<TableRow>> rows = new ArrayList<>();
  rows.add(wrapValue(new TableRow()));

  // First response is 403 non-{rate-limited, quota-exceeded}, second response has valid payload
  // but should not
  // be invoked.
  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(403).thenReturn(200);
  when(response.getContent())
      .thenReturn(toStream(errorWithReasonAndStatus("actually forbidden", 403)))
      .thenReturn(toStream(new TableDataInsertAllResponse()));

  DatasetServiceImpl dataService =
      new DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());
  dataService.insertAll(
      ref,
      rows,
      null,
      BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()),
      new MockSleeper(),
      InsertRetryPolicy.alwaysRetry(),
      null,
      null,
      false,
      false,
      false);
  verify(response, times(2)).getStatusCode();
  verify(response, times(2)).getContent();
  verify(response, times(2)).getContentType();
  expectedLogs.verifyInfo("BigQuery insertAll error, retrying:");
}
 

/** Tests that {@link DatasetServiceImpl#insertAll} retries quota exceeded attempts. */
@Test
public void testInsertQuotaExceededRetry() throws Exception {
  TableReference ref =
      new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
  List<ValueInSingleWindow<TableRow>> rows = new ArrayList<>();
  rows.add(wrapValue(new TableRow()));

  // First response is 403 quota exceeded, second response has valid payload.
  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(403).thenReturn(200);
  when(response.getContent())
      .thenReturn(toStream(errorWithReasonAndStatus("quotaExceeded", 403)))
      .thenReturn(toStream(new TableDataInsertAllResponse()));

  DatasetServiceImpl dataService =
      new DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());
  dataService.insertAll(
      ref,
      rows,
      null,
      BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()),
      new MockSleeper(),
      InsertRetryPolicy.alwaysRetry(),
      null,
      null,
      false,
      false,
      false);
  verify(response, times(2)).getStatusCode();
  verify(response, times(2)).getContent();
  verify(response, times(2)).getContentType();
  expectedLogs.verifyInfo("BigQuery insertAll error, retrying:");
}
 

/** Tests that {@link DatasetServiceImpl#insertAll} retries rate limited attempts. */
@Test
public void testInsertRateLimitRetry() throws Exception {
  TableReference ref =
      new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
  List<ValueInSingleWindow<TableRow>> rows = new ArrayList<>();
  rows.add(wrapValue(new TableRow()));

  // First response is 403 rate limited, second response has valid payload.
  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(403).thenReturn(200);
  when(response.getContent())
      .thenReturn(toStream(errorWithReasonAndStatus("rateLimitExceeded", 403)))
      .thenReturn(toStream(new TableDataInsertAllResponse()));

  DatasetServiceImpl dataService =
      new DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());
  dataService.insertAll(
      ref,
      rows,
      null,
      BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()),
      new MockSleeper(),
      InsertRetryPolicy.alwaysRetry(),
      null,
      null,
      false,
      false,
      false);
  verify(response, times(2)).getStatusCode();
  verify(response, times(2)).getContent();
  verify(response, times(2)).getContentType();
  expectedLogs.verifyInfo("BigQuery insertAll error, retrying:");
}
 

@Test
public void testMaxRetriesCountIsCumulative() throws IOException {
  MockSleeper sleeper = new MockSleeper();
  RetryInitializer initializer = new RetryInitializer(
      retryOnIOAndServiceUnavailableErrors(sleeper));

  final AtomicInteger counter = new AtomicInteger(0);
  MockLowLevelHttpRequest failingRequest = new MockLowLevelHttpRequest(){
    @Override
    public LowLevelHttpResponse execute() throws IOException {
      if (counter.getAndIncrement() < 2) {
        throw new IOException("test error");
      } else {
        return new MockLowLevelHttpResponse().setStatusCode(503).setZeroContent();
      }
    }
  };
  HttpRequest request = TestUtils.createRequest(failingRequest);
  initializer.initialize(request);
  final HttpUnsuccessfulResponseHandler retryHandler = request.getUnsuccessfulResponseHandler();

  try {
    request.execute();
    fail("No exception thrown for HTTP error");
  } catch (HttpResponseException e) {
    assertEquals(503, e.getStatusCode());
  }

  assertEquals(MAX_RETRIES, sleeper.getCount());
  assertEquals(MAX_RETRIES + 1, counter.get());
  assertSame(retryHandler, request.getUnsuccessfulResponseHandler());
}
 

@Test
public void testEnableRetry() throws IOException {
  RetryConfig retryConfig = retryOnIOAndServiceUnavailableErrors(new MockSleeper());
  RetryInitializer initializer = new RetryInitializer(retryConfig);
  HttpRequest request = TestUtils.createRequest();

  initializer.initialize(request);

  assertEquals(MAX_RETRIES, request.getNumberOfRetries());
  assertTrue(request.getUnsuccessfulResponseHandler() instanceof RetryHandlerDecorator);
  RetryUnsuccessfulResponseHandler retryHandler =
      ((RetryHandlerDecorator) request.getUnsuccessfulResponseHandler()).getRetryHandler();
  assertSame(retryConfig, retryHandler.getRetryConfig());
  assertTrue(request.getIOExceptionHandler() instanceof HttpBackOffIOExceptionHandler);
}
 

@Test
public void testBuilderWithAllSettings() {
  ImmutableList<Integer> statusCodes = ImmutableList.of(500, 503);
  Sleeper sleeper = new MockSleeper();
  RetryConfig config = RetryConfig.builder()
      .setMaxRetries(4)
      .setRetryStatusCodes(statusCodes)
      .setRetryOnIOExceptions(true)
      .setMaxIntervalMillis(5 * 60 * 1000)
      .setBackOffMultiplier(1.5)
      .setSleeper(sleeper)
      .build();

  assertEquals(2, config.getRetryStatusCodes().size());
  assertEquals(statusCodes.get(0), config.getRetryStatusCodes().get(0));
  assertEquals(statusCodes.get(1), config.getRetryStatusCodes().get(1));
  assertTrue(config.isRetryOnIOExceptions());
  assertEquals(4, config.getMaxRetries());
  assertEquals(5 * 60 * 1000, config.getMaxIntervalMillis());
  assertEquals(1.5, config.getBackOffMultiplier(), 0.01);
  assertSame(sleeper, config.getSleeper());

  ExponentialBackOff backOff = (ExponentialBackOff) config.newBackOff();
  assertEquals(500, backOff.getInitialIntervalMillis());
  assertEquals(5 * 60 * 1000, backOff.getMaxIntervalMillis());
  assertEquals(1.5, backOff.getMultiplier(), 0.01);
  assertEquals(0.0, backOff.getRandomizationFactor(), 0.01);
  assertNotSame(backOff, config.newBackOff());
}
 

/** Tests that {@link DatasetServiceImpl#insertAll} uses the supplied {@link ErrorContainer}. */
@Test
public void testSimpleErrorRetrieval() throws InterruptedException, IOException {
  TableReference ref =
      new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
  List<ValueInSingleWindow<TableRow>> rows =
      ImmutableList.of(
          wrapValue(new TableRow().set("a", 1)), wrapValue(new TableRow().set("b", 2)));

  final TableDataInsertAllResponse failures =
      new TableDataInsertAllResponse()
          .setInsertErrors(
              ImmutableList.of(
                  new InsertErrors()
                      .setIndex(0L)
                      .setErrors(ImmutableList.of(new ErrorProto().setReason("timeout"))),
                  new InsertErrors()
                      .setIndex(1L)
                      .setErrors(ImmutableList.of(new ErrorProto().setReason("invalid")))));

  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(200);
  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);

  when(response.getContent()).thenReturn(toStream(failures));

  DatasetServiceImpl dataService =
      new DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());

  List<ValueInSingleWindow<TableRow>> failedInserts = Lists.newArrayList();
  dataService.insertAll(
      ref,
      rows,
      null,
      BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()),
      new MockSleeper(),
      InsertRetryPolicy.neverRetry(),
      failedInserts,
      ErrorContainer.TABLE_ROW_ERROR_CONTAINER,
      false,
      false,
      false);

  assertThat(failedInserts, is(rows));
}
 

/** Tests that {@link DatasetServiceImpl#insertAll} uses the supplied {@link ErrorContainer}. */
@Test
public void testExtendedErrorRetrieval() throws InterruptedException, IOException {
  TableReference ref =
      new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
  List<ValueInSingleWindow<TableRow>> rows =
      ImmutableList.of(
          wrapValue(new TableRow().set("a", 1)), wrapValue(new TableRow().set("b", 2)));

  final TableDataInsertAllResponse failures =
      new TableDataInsertAllResponse()
          .setInsertErrors(
              ImmutableList.of(
                  new InsertErrors()
                      .setIndex(0L)
                      .setErrors(ImmutableList.of(new ErrorProto().setReason("timeout"))),
                  new InsertErrors()
                      .setIndex(1L)
                      .setErrors(ImmutableList.of(new ErrorProto().setReason("invalid")))));

  final List<ValueInSingleWindow<BigQueryInsertError>> expected =
      ImmutableList.of(
          wrapValue(
              new BigQueryInsertError(
                  rows.get(0).getValue(), failures.getInsertErrors().get(0), ref)),
          wrapValue(
              new BigQueryInsertError(
                  rows.get(1).getValue(), failures.getInsertErrors().get(1), ref)));

  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(200);
  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);

  when(response.getContent()).thenReturn(toStream(failures));

  DatasetServiceImpl dataService =
      new DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());

  List<ValueInSingleWindow<BigQueryInsertError>> failedInserts = Lists.newArrayList();
  dataService.insertAll(
      ref,
      rows,
      null,
      BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()),
      new MockSleeper(),
      InsertRetryPolicy.neverRetry(),
      failedInserts,
      ErrorContainer.BIG_QUERY_INSERT_ERROR_ERROR_CONTAINER,
      false,
      false,
      false);

  assertThat(failedInserts, is(expected));
}
 

/**
 * Tests that {@link DatasetServiceImpl#insertAll} respects the skipInvalidRows,
 * ignoreUnknownValues and ignoreInsertIds parameters.
 */
@Test
public void testSkipInvalidRowsIgnoreUnknownIgnoreInsertIdsValuesStreaming()
    throws InterruptedException, IOException {
  TableReference ref =
      new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
  List<ValueInSingleWindow<TableRow>> rows =
      ImmutableList.of(wrapValue(new TableRow()), wrapValue(new TableRow()));

  final TableDataInsertAllResponse allRowsSucceeded = new TableDataInsertAllResponse();

  // Return a 200 response each time
  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(200);
  when(response.getContent())
      .thenReturn(toStream(allRowsSucceeded))
      .thenReturn(toStream(allRowsSucceeded));

  DatasetServiceImpl dataService =
      new DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());

  // First, test with all flags disabled
  dataService.insertAll(
      ref,
      rows,
      null,
      BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()),
      new MockSleeper(),
      InsertRetryPolicy.neverRetry(),
      Lists.newArrayList(),
      ErrorContainer.TABLE_ROW_ERROR_CONTAINER,
      false,
      false,
      false);

  TableDataInsertAllRequest parsedRequest =
      fromString(request.getContentAsString(), TableDataInsertAllRequest.class);

  assertFalse(parsedRequest.getSkipInvalidRows());
  assertFalse(parsedRequest.getIgnoreUnknownValues());

  // Then with all enabled
  dataService.insertAll(
      ref,
      rows,
      null,
      BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()),
      new MockSleeper(),
      InsertRetryPolicy.neverRetry(),
      Lists.newArrayList(),
      ErrorContainer.TABLE_ROW_ERROR_CONTAINER,
      true,
      true,
      true);

  parsedRequest = fromString(request.getContentAsString(), TableDataInsertAllRequest.class);

  assertTrue(parsedRequest.getSkipInvalidRows());
  assertTrue(parsedRequest.getIgnoreUnknownValues());
  assertNull(parsedRequest.getRows().get(0).getInsertId());
  assertNull(parsedRequest.getRows().get(1).getInsertId());
}
 

/**
 * Tests that {@link DatasetServiceImpl#insertAll} uses the supplied {@link InsertRetryPolicy},
 * and returns the list of rows not retried.
 */
@Test
public void testInsertRetryPolicy() throws InterruptedException, IOException {
  TableReference ref =
      new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
  List<ValueInSingleWindow<TableRow>> rows =
      ImmutableList.of(wrapValue(new TableRow()), wrapValue(new TableRow()));

  // First time row0 fails with a retryable error, and row1 fails with a persistent error.
  final TableDataInsertAllResponse firstFailure =
      new TableDataInsertAllResponse()
          .setInsertErrors(
              ImmutableList.of(
                  new InsertErrors()
                      .setIndex(0L)
                      .setErrors(ImmutableList.of(new ErrorProto().setReason("timeout"))),
                  new InsertErrors()
                      .setIndex(1L)
                      .setErrors(ImmutableList.of(new ErrorProto().setReason("invalid")))));

  // Second time there is only one row, which fails with a retryable error.
  final TableDataInsertAllResponse secondFialure =
      new TableDataInsertAllResponse()
          .setInsertErrors(
              ImmutableList.of(
                  new InsertErrors()
                      .setIndex(0L)
                      .setErrors(ImmutableList.of(new ErrorProto().setReason("timeout")))));

  // On the final attempt, no failures are returned.
  final TableDataInsertAllResponse allRowsSucceeded = new TableDataInsertAllResponse();

  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  // Always return 200.
  when(response.getStatusCode()).thenReturn(200);
  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(200).thenReturn(200);

  // First fail
  when(response.getContent())
      .thenReturn(toStream(firstFailure))
      .thenReturn(toStream(secondFialure))
      .thenReturn(toStream(allRowsSucceeded));

  DatasetServiceImpl dataService =
      new DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());

  List<ValueInSingleWindow<TableRow>> failedInserts = Lists.newArrayList();
  dataService.insertAll(
      ref,
      rows,
      null,
      BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()),
      new MockSleeper(),
      InsertRetryPolicy.retryTransientErrors(),
      failedInserts,
      ErrorContainer.TABLE_ROW_ERROR_CONTAINER,
      false,
      false,
      false);
  assertEquals(1, failedInserts.size());
  expectedLogs.verifyInfo("Retrying 1 failed inserts to BigQuery");
}
 

/** Tests that {@link DatasetServiceImpl#insertAll} fails gracefully when persistent issues. */
@Test
public void testInsertFailsGracefully() throws Exception {
  TableReference ref =
      new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
  List<ValueInSingleWindow<TableRow>> rows =
      ImmutableList.of(wrapValue(new TableRow()), wrapValue(new TableRow()));

  final TableDataInsertAllResponse row1Failed =
      new TableDataInsertAllResponse()
          .setInsertErrors(ImmutableList.of(new InsertErrors().setIndex(1L)));

  final TableDataInsertAllResponse row0Failed =
      new TableDataInsertAllResponse()
          .setInsertErrors(ImmutableList.of(new InsertErrors().setIndex(0L)));

  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  // Always return 200.
  when(response.getStatusCode()).thenReturn(200);
  // Return row 1 failing, then we retry row 1 as row 0, and row 0 persistently fails.
  when(response.getContent())
      .thenReturn(toStream(row1Failed))
      .thenAnswer(invocation -> toStream(row0Failed));

  DatasetServiceImpl dataService =
      new DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());

  // Expect it to fail.
  try {
    dataService.insertAll(
        ref,
        rows,
        null,
        BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()),
        new MockSleeper(),
        InsertRetryPolicy.alwaysRetry(),
        null,
        null,
        false,
        false,
        false);
    fail();
  } catch (IOException e) {
    assertThat(e, instanceOf(IOException.class));
    assertThat(e.getMessage(), containsString("Insert failed:"));
    assertThat(e.getMessage(), containsString("[{\"index\":0}]"));
  }

  // Verify the exact number of retries as well as log messages.
  verify(response, times(4)).getStatusCode();
  verify(response, times(4)).getContent();
  verify(response, times(4)).getContentType();
  expectedLogs.verifyInfo("Retrying 1 failed inserts to BigQuery");
}
 

/** Tests that {@link DatasetServiceImpl#insertAll} retries selected rows on failure. */
@Test
public void testInsertRetrySelectRows() throws Exception {
  TableReference ref =
      new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
  List<ValueInSingleWindow<TableRow>> rows =
      ImmutableList.of(
          wrapValue(new TableRow().set("row", "a")), wrapValue(new TableRow().set("row", "b")));
  List<String> insertIds = ImmutableList.of("a", "b");

  final TableDataInsertAllResponse bFailed =
      new TableDataInsertAllResponse()
          .setInsertErrors(
              ImmutableList.of(
                  new InsertErrors().setIndex(1L).setErrors(ImmutableList.of(new ErrorProto()))));

  final TableDataInsertAllResponse allRowsSucceeded = new TableDataInsertAllResponse();

  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(200).thenReturn(200);
  when(response.getContent())
      .thenReturn(toStream(bFailed))
      .thenReturn(toStream(allRowsSucceeded));

  DatasetServiceImpl dataService =
      new DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());
  dataService.insertAll(
      ref,
      rows,
      insertIds,
      BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()),
      new MockSleeper(),
      InsertRetryPolicy.alwaysRetry(),
      null,
      null,
      false,
      false,
      false);
  verify(response, times(2)).getStatusCode();
  verify(response, times(2)).getContent();
  verify(response, times(2)).getContentType();
}