com.google.api.client.googleapis.batch.BatchRequest Java Examples

@Test
public void testBatchRequestIOException() throws Exception {
  BatchRequestService batchService = setupService();
  batchService.startAsync().awaitRunning();
  assertTrue(batchService.isRunning());
  BatchRequest batchRequest = getMockBatchRequest();
  when(batchRequestHelper.createBatch(any())).thenReturn(batchRequest);
  AsyncRequest<GenericJson> requestToBatch =
      new AsyncRequest<GenericJson>(testRequest, retryPolicy, operationStats);
  GoogleJsonError error = new GoogleJsonError();
  error.setCode(403);
  error.setMessage("unauthorized");
  doAnswer(
          invocation -> {
            throw new IOException();
          })
      .when(batchRequestHelper)
      .executeBatchRequest(batchRequest);

  batchService.add(requestToBatch);
  batchService.flush();
  batchService.stopAsync().awaitTerminated();
  validateFailedResult(requestToBatch.getFuture());
  assertFalse(batchService.isRunning());
}
 

@Test
public void testMakeRemoveBatches() throws IOException {
  GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil();

  // Small number of files fits in 1 batch
  List<BatchRequest> batches = gcsUtil.makeRemoveBatches(makeStrings("s", 3));
  assertThat(batches.size(), equalTo(1));
  assertThat(sumBatchSizes(batches), equalTo(3));

  // 1 batch of files fits in 1 batch
  batches = gcsUtil.makeRemoveBatches(makeStrings("s", 100));
  assertThat(batches.size(), equalTo(1));
  assertThat(sumBatchSizes(batches), equalTo(100));

  // A little more than 5 batches of files fits in 6 batches
  batches = gcsUtil.makeRemoveBatches(makeStrings("s", 501));
  assertThat(batches.size(), equalTo(6));
  assertThat(sumBatchSizes(batches), equalTo(501));
}
 

@Test
public void testMakeGetBatches() throws IOException {
  GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil();

  // Small number of files fits in 1 batch
  List<StorageObjectOrIOException[]> results = Lists.newArrayList();
  List<BatchRequest> batches = gcsUtil.makeGetBatches(makeGcsPaths("s", 3), results);
  assertThat(batches.size(), equalTo(1));
  assertThat(sumBatchSizes(batches), equalTo(3));
  assertEquals(3, results.size());

  // 1 batch of files fits in 1 batch
  results = Lists.newArrayList();
  batches = gcsUtil.makeGetBatches(makeGcsPaths("s", 100), results);
  assertThat(batches.size(), equalTo(1));
  assertThat(sumBatchSizes(batches), equalTo(100));
  assertEquals(100, results.size());

  // A little more than 5 batches of files fits in 6 batches
  results = Lists.newArrayList();
  batches = gcsUtil.makeGetBatches(makeGcsPaths("s", 501), results);
  assertThat(batches.size(), equalTo(6));
  assertThat(sumBatchSizes(batches), equalTo(501));
  assertEquals(501, results.size());
}
 

@Test
public void testMakeCopyBatches() throws IOException {
  GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil();

  // Small number of files fits in 1 batch
  List<BatchRequest> batches =
      gcsUtil.makeCopyBatches(gcsUtil.makeRewriteOps(makeStrings("s", 3), makeStrings("d", 3)));
  assertThat(batches.size(), equalTo(1));
  assertThat(sumBatchSizes(batches), equalTo(3));

  // 1 batch of files fits in 1 batch
  batches =
      gcsUtil.makeCopyBatches(
          gcsUtil.makeRewriteOps(makeStrings("s", 100), makeStrings("d", 100)));
  assertThat(batches.size(), equalTo(1));
  assertThat(sumBatchSizes(batches), equalTo(100));

  // A little more than 5 batches of files fits in 6 batches
  batches =
      gcsUtil.makeCopyBatches(
          gcsUtil.makeRewriteOps(makeStrings("s", 501), makeStrings("d", 501)));
  assertThat(batches.size(), equalTo(6));
  assertThat(sumBatchSizes(batches), equalTo(501));
}
 

List<BatchRequest> makeCopyBatches(LinkedList<RewriteOp> rewrites) throws IOException {
  List<BatchRequest> batches = new ArrayList<>();
  BatchRequest batch = createBatchRequest();
  Iterator<RewriteOp> it = rewrites.iterator();
  while (it.hasNext()) {
    RewriteOp rewrite = it.next();
    if (!rewrite.getReadyToEnqueue()) {
      it.remove();
      continue;
    }
    rewrite.enqueue(batch);

    if (batch.size() >= MAX_REQUESTS_PER_BATCH) {
      batches.add(batch);
      batch = createBatchRequest();
    }
  }
  if (batch.size() > 0) {
    batches.add(batch);
  }
  return batches;
}
 

@Override
public void performSearchAnalyticsQueryInBatch(List<ProducerJob> jobs, List<ArrayList<ApiDimensionFilter>> filterList,
    List<JsonBatchCallback<SearchAnalyticsQueryResponse>> callbackList, List<Dimension> requestedDimensions,
    int rowLimit) throws IOException {
  BatchRequest batchRequest = _client.createBatch();

  for (int i = 0; i < jobs.size(); ++i) {
    ProducerJob job = jobs.get(i);
    ArrayList<ApiDimensionFilter> filters = filterList.get(i);
    JsonBatchCallback<SearchAnalyticsQueryResponse> callback = callbackList.get(i);
    _client.createSearchAnalyticsQuery(_siteProperty, job.getStartDate(), job.getEndDate(), requestedDimensions,
        GoogleWebmasterFilter.andGroupFilters(filters), rowLimit, 0).queue(batchRequest, callback);
  }

  batchRequest.execute();
}
 

@Test
public void testFailedResult() throws Exception {
  BatchRequestService batchService = setupService();
  batchService.startAsync().awaitRunning();
  assertTrue(batchService.isRunning());
  BatchRequest batchRequest = getMockBatchRequest();
  when(batchRequestHelper.createBatch(any())).thenReturn(batchRequest);
  AsyncRequest<GenericJson> requestToBatch =
      new AsyncRequest<GenericJson>(testRequest, retryPolicy, operationStats);
  GoogleJsonError error = new GoogleJsonError();
  error.setCode(403);
  error.setMessage("unauthorized");
  doAnswer(
          invocation -> {
            requestToBatch.getCallback().onFailure(error, new HttpHeaders());
            return null;
          })
      .when(batchRequestHelper)
      .executeBatchRequest(batchRequest);

  batchService.add(requestToBatch);
  batchService.flush();
  batchService.stopAsync().awaitTerminated();
  validateFailedResult(requestToBatch.getFuture());
  assertFalse(batchService.isRunning());
}
 

@Test
public void testFlushOnShutdown() throws Exception {
  BatchRequestService batchService = setupService();
  batchService.startAsync().awaitRunning();
  assertTrue(batchService.isRunning());
  BatchRequest batchRequest = getMockBatchRequest();
  when(batchRequestHelper.createBatch(any())).thenReturn(batchRequest);
  AsyncRequest<GenericJson> requestToBatch =
      new AsyncRequest<GenericJson>(testRequest, retryPolicy, operationStats);
  doAnswer(
          invocation -> {
            requestToBatch.getCallback().onStart();
            requestToBatch.getCallback().onSuccess(new GenericJson(), new HttpHeaders());
            return null;
          })
      .when(batchRequestHelper)
      .executeBatchRequest(batchRequest);

  batchService.add(requestToBatch);
  batchService.stopAsync().awaitTerminated();
  assertEquals(new GenericJson(), requestToBatch.getFuture().get());
  assertFalse(batchService.isRunning());
}
 

@Test
public void testFlush() throws Exception {
  BatchRequestService batchService = setupService();
  batchService.startAsync().awaitRunning();
  assertTrue(batchService.isRunning());
  BatchRequest batchRequest = getMockBatchRequest();
  when(batchRequestHelper.createBatch(any())).thenReturn(batchRequest);
  AsyncRequest<GenericJson> requestToBatch =
      new AsyncRequest<GenericJson>(testRequest, retryPolicy, operationStats);
  doAnswer(
          invocation -> {
            requestToBatch.getCallback().onStart();
            requestToBatch.getCallback().onSuccess(new GenericJson(), new HttpHeaders());
            return null;
          })
      .when(batchRequestHelper)
      .executeBatchRequest(batchRequest);

  batchService.add(requestToBatch);
  batchService.flush();
  batchService.stopAsync().awaitTerminated();
  assertEquals(new GenericJson(), requestToBatch.getFuture().get());
  assertFalse(batchService.isRunning());
}
 

private BatchRequest newBatchRequest(
    List<Message> messages, boolean dryRun, MessagingBatchCallback callback) throws IOException {

  BatchRequest batch = new BatchRequest(
      requestFactory.getTransport(), getBatchRequestInitializer());
  batch.setBatchUrl(new GenericUrl(FCM_BATCH_URL));

  final JsonObjectParser jsonParser = new JsonObjectParser(this.jsonFactory);
  final GenericUrl sendUrl = new GenericUrl(fcmSendUrl);
  for (Message message : messages) {
    // Using a separate request factory without authorization is faster for large batches.
    // A simple performance test showed a 400-500ms speed up for batches of 1000 messages.
    HttpRequest request = childRequestFactory.buildPostRequest(
        sendUrl,
        new JsonHttpContent(jsonFactory, message.wrapForTransport(dryRun)));
    request.setParser(jsonParser);
    setCommonFcmHeaders(request.getHeaders());
    batch.queue(
        request, MessagingServiceResponse.class, MessagingServiceErrorResponse.class, callback);
  }
  return batch;
}
 

private void flushPendingRequests() throws IOException {
  if (pendingRequests.isEmpty()) {
    return;
  }
  BatchRequest batch = gcs.batch(requestInitializer);
  while (batch.size() < maxRequestsPerBatch && !pendingRequests.isEmpty()) {
    // enqueue request at head
    pendingRequests.remove().enqueue(batch);
  }
  responseFutures.add(
      requestsExecutor.submit(
          () -> {
            batch.execute();
            return null;
          }));
}
 

@Test
@SuppressWarnings("unchecked")
public void testNonRetryableErrorCode() throws Exception {
  int httpErrorCode = 65535;
  when(retryPolicy.isRetryableStatusCode(httpErrorCode)).thenReturn(false);
  BatchRequestService batchService = setupService();
  batchService.startAsync().awaitRunning();
  assertTrue(batchService.isRunning());
  BatchRequest batchRequest = getMockBatchRequest();
  when(batchRequestHelper.createBatch(any())).thenReturn(batchRequest);
  AsyncRequest<GenericJson> requestToBatch =
      new AsyncRequest<GenericJson>(testRequest, retryPolicy, operationStats);
  assertEquals(0, requestToBatch.getRetries());
  assertEquals(Status.NEW, requestToBatch.getStatus());
  GoogleJsonError error = new GoogleJsonError();
  error.setCode(httpErrorCode);
  error.setMessage("Unknown error code");

  doAnswer(
          i -> {
            requestToBatch.getCallback().onFailure(error, new HttpHeaders());
            return null;
          })
      .when(batchRequestHelper)
      .executeBatchRequest(batchRequest);

  batchService.add(requestToBatch);
  batchService.flush();
  batchService.stopAsync().awaitTerminated();
  verify(retryPolicy).isRetryableStatusCode(httpErrorCode);
  validateFailedResult(requestToBatch.getFuture());
  assertEquals(Status.FAILED, requestToBatch.getStatus());
  assertEquals(0, requestToBatch.getRetries());
  assertFalse(batchService.isRunning());
}
 

private BatchResponse sendBatchRequest(
    List<Message> messages, boolean dryRun) throws IOException {

  MessagingBatchCallback callback = new MessagingBatchCallback();
  BatchRequest batch = newBatchRequest(messages, dryRun, callback);
  batch.execute();
  return new BatchResponseImpl(callback.getResponses());
}
 

List<BatchRequest> makeRemoveBatches(Collection<String> filenames) throws IOException {
  List<BatchRequest> batches = new ArrayList<>();
  for (List<String> filesToDelete :
      Lists.partition(Lists.newArrayList(filenames), MAX_REQUESTS_PER_BATCH)) {
    BatchRequest batch = createBatchRequest();
    for (String file : filesToDelete) {
      enqueueDelete(GcsPath.fromUri(file), batch);
    }
    batches.add(batch);
  }
  return batches;
}
 

private static int sumBatchSizes(List<BatchRequest> batches) {
  int ret = 0;
  for (BatchRequest b : batches) {
    ret += b.size();
    assertThat(b.size(), greaterThan(0));
  }
  return ret;
}
 

/**
 * Makes get {@link BatchRequest BatchRequests}.
 *
 * @param paths {@link GcsPath GcsPaths}.
 * @param results mutable {@link List} for return values.
 * @return {@link BatchRequest BatchRequests} to execute.
 * @throws IOException
 */
@VisibleForTesting
List<BatchRequest> makeGetBatches(
    Collection<GcsPath> paths, List<StorageObjectOrIOException[]> results) throws IOException {
  List<BatchRequest> batches = new ArrayList<>();
  for (List<GcsPath> filesToGet :
      Lists.partition(Lists.newArrayList(paths), MAX_REQUESTS_PER_BATCH)) {
    BatchRequest batch = createBatchRequest();
    for (GcsPath path : filesToGet) {
      results.add(enqueueGetFileSize(path, batch));
    }
    batches.add(batch);
  }
  return batches;
}
 

public void enqueue(BatchRequest batch) throws IOException {
  if (!readyToEnqueue) {
    throw new IOException(
        String.format(
            "Invalid state for Rewrite, from=%s, to=%s, readyToEnqueue=%s",
            from, to, readyToEnqueue));
  }
  rewriteRequest.queue(batch, this);
  readyToEnqueue = false;
}
 

@Test
@SuppressWarnings("unchecked")
public void testNonRetryableException() throws Exception {
  when(retryPolicy.isRetryableStatusCode(0)).thenReturn(false);
  BatchRequestService batchService = setupService();
  batchService.startAsync().awaitRunning();
  assertTrue(batchService.isRunning());
  BatchRequest batchRequest = getMockBatchRequest();
  when(batchRequestHelper.createBatch(any())).thenReturn(batchRequest);
  AsyncRequest<GenericJson> requestToBatch =
      new AsyncRequest<GenericJson>(testRequest, retryPolicy, operationStats);
  assertEquals(0, requestToBatch.getRetries());
  assertEquals(Status.NEW, requestToBatch.getStatus());

  doAnswer(
          i -> {
            throw new IOException("Non-retryable exception");
          })
      .when(batchRequestHelper)
      .executeBatchRequest(batchRequest);

  batchService.add(requestToBatch);
  batchService.flush();
  batchService.stopAsync().awaitTerminated();
  verify(retryPolicy).isRetryableStatusCode(0);
  validateFailedResult(requestToBatch.getFuture());
  assertEquals(Status.FAILED, requestToBatch.getStatus());
  assertEquals(0, requestToBatch.getRetries());
  assertFalse(batchService.isRunning());
}
 

@Test
@SuppressWarnings("unchecked")
public void testAllRetriesFailed() throws Exception {
  int httpErrorCode = 503;
  when(retryPolicy.isRetryableStatusCode(httpErrorCode)).thenReturn(true);
  int retries = 3;
  when(retryPolicy.getMaxRetryLimit()).thenReturn(retries);
  BatchRequestService batchService = setupService();
  batchService.startAsync().awaitRunning();
  assertTrue(batchService.isRunning());
  BatchRequest batchRequest = getMockBatchRequest();
  when(batchRequestHelper.createBatch(any())).thenReturn(batchRequest);
  AsyncRequest<GenericJson> requestToBatch =
      new AsyncRequest<GenericJson>(testRequest, retryPolicy, operationStats);
  assertEquals(0, requestToBatch.getRetries());
  assertEquals(Status.NEW, requestToBatch.getStatus());
  GoogleJsonError error = new GoogleJsonError();
  error.setCode(httpErrorCode);
  error.setMessage("Service Unavailable");

  doAnswer(
          i -> {
            requestToBatch.getCallback().onFailure(error, new HttpHeaders());
            return null;
          })
      .when(batchRequestHelper)
      .executeBatchRequest(batchRequest);

  batchService.add(requestToBatch);
  batchService.flush();
  batchService.stopAsync().awaitTerminated();
  verify(retryPolicy, times(retries + 1)).isRetryableStatusCode(httpErrorCode);
  verify(retryPolicy, times(retries + 1)).getMaxRetryLimit();
  validateFailedResult(requestToBatch.getFuture());
  assertEquals(Status.FAILED, requestToBatch.getStatus());
  assertEquals(retries, requestToBatch.getRetries());
  assertFalse(batchService.isRunning());
}
 

@Test
public void testBatchRequestSocketTimeoutException() throws Exception {
  BatchRequestService batchService = setupService();
  batchService.startAsync().awaitRunning();
  assertTrue(batchService.isRunning());
  BatchRequest batchRequest = getMockBatchRequest();
  when(batchRequestHelper.createBatch(any())).thenReturn(batchRequest);
  when(retryPolicy.getMaxRetryLimit()).thenReturn(1);
  when(retryPolicy.isRetryableStatusCode(504)).thenReturn(true);
  AsyncRequest<GenericJson> requestToBatch =
      new AsyncRequest<GenericJson>(testRequest, retryPolicy, operationStats);
  AtomicInteger counter = new AtomicInteger();
  GenericJson successfulResult = new GenericJson();
  doAnswer(
          invocation -> {
            if (counter.incrementAndGet() == 1) {
              throw new SocketTimeoutException();
            }
            requestToBatch.getCallback().onStart();
            requestToBatch.getCallback().onSuccess(successfulResult, new HttpHeaders());
            return null;
          })
      .when(batchRequestHelper)
      .executeBatchRequest(batchRequest);

  batchService.add(requestToBatch);
  batchService.flush();
  batchService.stopAsync().awaitTerminated();
  assertEquals(successfulResult, requestToBatch.getFuture().get());
  assertEquals(Status.COMPLETED, requestToBatch.getStatus());
  assertEquals(1, requestToBatch.getRetries());
  assertFalse(batchService.isRunning());
}
 

<T> void queue(
    BatchRequest batchRequest,
    AbstractGoogleJsonClientRequest<T> request,
    JsonBatchCallback<T> jsonCallback)
    throws IOException {
  request.queue(batchRequest, jsonCallback);
}
 

@SuppressWarnings("unchecked")
private void execute(List<AsyncRequest<?>> toBeProcessed) throws IOException {
  List<? extends EventStartCallback> eventCallbackList =
      ImmutableList.copyOf(
          toBeProcessed.stream().map(r -> r.getCallback()).collect(Collectors.toList()));
  BatchRequest batchRequest =
      batchRequestHelper.createBatch(
          new EventLoggingRequestInitializer(batchRequestInitializer, eventCallbackList));
  for (@SuppressWarnings("rawtypes") AsyncRequest req : toBeProcessed) {
    batchRequestHelper.queue(batchRequest, req.getRequest(), req.getCallback());
  }
  batchRequestHelper.executeBatchRequest(batchRequest);
}
 

BatchRequest createBatch(HttpRequestInitializer requestInitializer) {
  return service.batch(requestInitializer);
}
 

private BatchRequest createBatchRequest() {
  return storageClient.batch(httpRequestInitializer);
}
 

private BatchRequest buildBatchRequest(Compute compute) {
  return compute.batch(
      (HttpRequest request) -> {
        request.getHeaders().setUserAgent("halyard " + halyardVersion);
      });
}
 

@Override
public BatchRequest createBatch() {
  return _service.batch();
}
 

@Test
public void batchPolicy_fromConfiguration_maxBatchDelay_custom() throws Exception {
  Integer batchDelay = 3;
  Properties config = new Properties();
  config.setProperty("batch.maxBatchDelaySeconds", String.valueOf(batchDelay));
  // Don't flush on shutdown; the test is checking for the scheduled flush.
  config.setProperty("batch.flushOnShutdown", "false");
  setupConfig.initConfig(config);
  BatchPolicy batchPolicy = BatchPolicy.fromConfiguration();
  assertEquals(batchDelay.intValue(), batchPolicy.getMaxBatchDelay());

  // Use a real factory so we get a working object for the ScheduledExecutorService
  ExecutorFactory executorFactory = spy(new BatchRequestService.ExecutorFactoryImpl());
  ExecutorService batchExecutor = Mockito.mock(ExecutorService.class);
  when(executorFactory.getExecutor()).thenReturn(batchExecutor);
  doAnswer(
          invocation -> {
            ((SnapshotRunnable) invocation.getArgument(0)).run();
            return null;
          })
      .when(batchExecutor)
      .execute(isA(SnapshotRunnable.class));
  AsyncRequest<GenericJson> requestToBatch =
      new AsyncRequest<GenericJson>(testRequest, retryPolicy, operationStats);
  BatchRequest batchRequest = getMockBatchRequest();
  when(batchRequestHelper.createBatch(any())).thenReturn(batchRequest);
  doAnswer(
          invocation -> {
            requestToBatch.getCallback().onStart();
            requestToBatch.getCallback().onSuccess(new GenericJson(), new HttpHeaders());
            return null;
          })
      .when(batchRequestHelper)
      .executeBatchRequest(batchRequest);

  BatchRequestService batchService =
      new BatchRequestService.Builder(service)
          .setExecutorFactory(executorFactory)
          .setBatchRequestHelper(batchRequestHelper)
          .setGoogleCredential(credential)
          .setBatchPolicy(batchPolicy)
          .build();
  batchService.startAsync().awaitRunning();
  batchService.add(requestToBatch);
  // Wait out the configured delay.
  Thread.sleep((batchDelay + 1) * 1000);
  batchService.stopAsync().awaitTerminated();

  // If the flush doesn't happen, the request will be cancelled and get() will throw an
  // exception.
  assertEquals(new GenericJson(), requestToBatch.getFuture().get());
  verify(batchExecutor).execute(isA(SnapshotRunnable.class));
}
 

@Test
public void batchPolicy_fromConfiguration_maxBatchDelay_default() throws Exception {
  Properties config = new Properties();
  // Don't flush on shutdown; the test is checking for the scheduled flush.
  config.setProperty("batch.flushOnShutdown", "false");
  setupConfig.initConfig(config);
  BatchPolicy batchPolicy = BatchPolicy.fromConfiguration();
  Integer batchDelay = batchPolicy.getMaxBatchDelay();
  // Check for a change in the default that will increase test run time.
  if (batchDelay > 20) {
    fail("Running this test with default batchDelay=" + batchDelay
        + " will make the test run time too long; remove this test?");
  }
  // Use a real factory so we get a working object for the ScheduledExecutorService
  ExecutorFactory executorFactory = spy(new BatchRequestService.ExecutorFactoryImpl());
  ExecutorService batchExecutor = Mockito.mock(ExecutorService.class);
  when(executorFactory.getExecutor()).thenReturn(batchExecutor);
  doAnswer(
          invocation -> {
            ((SnapshotRunnable) invocation.getArgument(0)).run();
            return null;
          })
      .when(batchExecutor)
      .execute(isA(SnapshotRunnable.class));
  AsyncRequest<GenericJson> requestToBatch =
      new AsyncRequest<GenericJson>(testRequest, retryPolicy, operationStats);
  BatchRequest batchRequest = getMockBatchRequest();
  when(batchRequestHelper.createBatch(any())).thenReturn(batchRequest);
  doAnswer(
          invocation -> {
            requestToBatch.getCallback().onStart();
            requestToBatch.getCallback().onSuccess(new GenericJson(), new HttpHeaders());
            return null;
          })
      .when(batchRequestHelper)
      .executeBatchRequest(batchRequest);

  BatchRequestService batchService =
      new BatchRequestService.Builder(service)
          .setExecutorFactory(executorFactory)
          .setBatchRequestHelper(batchRequestHelper)
          .setGoogleCredential(credential)
          .setBatchPolicy(batchPolicy)
          .build();
  batchService.startAsync().awaitRunning();
  batchService.add(requestToBatch);
  // Wait out the configured delay.
  Thread.sleep((batchDelay + 1) * 1000);
  batchService.stopAsync().awaitTerminated();

  // If the flush doesn't happen, the request will be cancelled and get() will throw an
  // exception.
  assertEquals(new GenericJson(), requestToBatch.getFuture().get());
  verify(batchExecutor).execute(isA(SnapshotRunnable.class));
}
 

@Test
@SuppressWarnings("unchecked")
public void testFailedRequestSucceedsOnRetry() throws Exception {
  int httpErrorCode = 503;
  when(retryPolicy.isRetryableStatusCode(httpErrorCode)).thenReturn(true);
  when(retryPolicy.getMaxRetryLimit()).thenReturn(1);
  BatchRequestService batchService = setupService();
  batchService.startAsync().awaitRunning();
  assertTrue(batchService.isRunning());
  BatchRequest batchRequest = getMockBatchRequest();
  when(batchRequestHelper.createBatch(any())).thenReturn(batchRequest);
  AsyncRequest<GenericJson> requestToBatch =
      new AsyncRequest<GenericJson>(testRequest, retryPolicy, operationStats);
  GoogleJsonError error = new GoogleJsonError();
  error.setCode(httpErrorCode);
  error.setMessage("Service Unavailable");
  GenericJson successfulResult = new GenericJson();
  doAnswer(
          i -> {
            if (requestToBatch.getRetries() >= 1) {
              requestToBatch.getCallback().onStart();
              requestToBatch.getCallback().onSuccess(successfulResult, new HttpHeaders());
            } else {
              requestToBatch.getCallback().onFailure(error, new HttpHeaders());
            }
            return null;
          })
      .when(batchRequestHelper)
      .executeBatchRequest(batchRequest);

  batchService.add(requestToBatch);
  batchService.flush();
  batchService.stopAsync().awaitTerminated();
  verify(retryPolicy).isRetryableStatusCode(httpErrorCode);
  verify(retryPolicy).getMaxRetryLimit();
  verify(backOff).nextBackOffMillis();
  assertEquals(successfulResult, requestToBatch.getFuture().get());
  assertEquals(Status.COMPLETED, requestToBatch.getStatus());
  assertEquals(1, requestToBatch.getRetries());
  assertFalse(batchService.isRunning());
}
 

@Test
@SuppressWarnings("unchecked")
public void testRetryableGoogleJSONException() throws Exception {
  int httpErrorCode = 503;
  String errorMessage = "Service Unavailable";
  when(retryPolicy.isRetryableStatusCode(httpErrorCode)).thenReturn(true);
  when(retryPolicy.getMaxRetryLimit()).thenReturn(3);
  BatchRequestService batchService = setupService();
  batchService.startAsync().awaitRunning();
  assertTrue(batchService.isRunning());
  BatchRequest batchRequest = getMockBatchRequest();
  when(batchRequestHelper.createBatch(any())).thenReturn(batchRequest);
  AsyncRequest<GenericJson> requestToBatch =
      new AsyncRequest<GenericJson>(testRequest, retryPolicy, operationStats);
  assertEquals(0, requestToBatch.getRetries());
  assertEquals(Status.NEW, requestToBatch.getStatus());
  GoogleJsonError error = new GoogleJsonError();
  error.setCode(httpErrorCode);
  error.setMessage(errorMessage);
  GoogleJsonResponseException exception =
      new GoogleJsonResponseException(
          new HttpResponseException.Builder(httpErrorCode, errorMessage, new HttpHeaders()),
          error);
  GenericJson successfulResult = new GenericJson();
  doAnswer(
          i -> {
            if (requestToBatch.getRetries() >= 2) {
              requestToBatch.getCallback().onStart();
              requestToBatch.getCallback().onSuccess(successfulResult, new HttpHeaders());
              return null;
            } else {
              throw exception;
            }
          })
      .when(batchRequestHelper)
      .executeBatchRequest(batchRequest);

  batchService.add(requestToBatch);
  batchService.flush();
  batchService.stopAsync().awaitTerminated();
  verify(retryPolicy, times(2)).isRetryableStatusCode(httpErrorCode);
  verify(retryPolicy, times(2)).getMaxRetryLimit();
  verify(backOff, times(2)).nextBackOffMillis();
  assertEquals(successfulResult, requestToBatch.getFuture().get());
  assertEquals(Status.COMPLETED, requestToBatch.getStatus());
  assertEquals(2, requestToBatch.getRetries());
  assertFalse(batchService.isRunning());
}