com.google.api.client.json.Json Java Examples

@Override
public LowLevelHttpRequest buildRequest(String method, String url) {
  return new MockLowLevelHttpRequest(url) {
    @Override
    public LowLevelHttpResponse execute() throws IOException {
      MockLowLevelHttpResponse response = new MockLowLevelHttpResponse();
      response.setContentType(Json.MEDIA_TYPE);
      IdTokenResponse json = new IdTokenResponse();
      json.setAccessToken("abc");
      json.setRefreshToken("def");
      json.setExpiresInSeconds(3600L);
      json.setIdToken(JWT_ENCODED_CONTENT);
      response.setContent(JSON_FACTORY.toString(json));
      return response;
    }
  };
}
 

public HttpResponse updateUserForHttpResponse(String username, java.io.InputStream body, String mediaType) throws IOException {
    // verify the required parameter 'username' is set
        if (username == null) {
        throw new IllegalArgumentException("Missing the required parameter 'username' when calling updateUser");
        }// verify the required parameter 'body' is set
        if (body == null) {
        throw new IllegalArgumentException("Missing the required parameter 'body' when calling updateUser");
        }
            // create a map of path variables
            final Map<String, Object> uriVariables = new HashMap<String, Object>();
                uriVariables.put("username", username);
        UriBuilder uriBuilder = UriBuilder.fromUri(apiClient.getBasePath() + "/user/{username}");

        String localVarUrl = uriBuilder.buildFromMap(uriVariables).toString();
        GenericUrl genericUrl = new GenericUrl(localVarUrl);

        HttpContent content = body == null ?
          apiClient.new JacksonJsonHttpContent(null) :
          new InputStreamContent(mediaType == null ? Json.MEDIA_TYPE : mediaType, body);
        return apiClient.getHttpRequestFactory().buildRequest(HttpMethods.PUT, genericUrl, content).execute();
}
 

@Override
public MockLowLevelHttpResponse execute() throws IOException {
  isExecuted.set(true);
  String requestContent = this.getContentAsString();
  T incomingRequest =
      JSON_PARSER.parseAndClose(new StringReader(requestContent), expectedRequestClass);
  boolean isExpected = expected.equals(incomingRequest);
  if (!isExpected) {
    hasError.set(true);
  }
  assertEquals(expected, incomingRequest);
  MockLowLevelHttpResponse response = new MockLowLevelHttpResponse();
  response
  .setStatusCode(HTTP_OK)
  .setContentType(Json.MEDIA_TYPE)
  .setContent(JSON_FACTORY.toString(expectedResponse));
  return response;
}
 

public static LowLevelHttpRequest mockResponseWithStatus(
    String responseContent, int status, MockRequestAssertion requestValidator) {
  return new MockLowLevelHttpRequest() {
    @Override
    public LowLevelHttpResponse execute() throws IOException {
      assertWithMessage(
              String.format(
                  "Request <%s> did not match predicate: '%s'",
                  this.getContentAsString(), requestValidator))
          .that(requestValidator.test(this.getContentAsString()))
          .isTrue();
      // Responses contain a IntputStream for content. Cannot be reused between for two
      // consecutive calls. We create new ones per call here.
      return new MockLowLevelHttpResponse()
          .setContentType(Json.MEDIA_TYPE)
          .setContent(responseContent)
          .setStatusCode(status);
    }
  };
}
 

@Test
public void simple() throws ServletException, IOException {
  HttpServletRequest request = mock(HttpServletRequest.class);
  HttpServletResponse response = mock(HttpServletResponse.class);
  BufferedReader reader = records(
      ImmutableMap.of(
          "time", new Date(TIME),
          "host", "localhost",
          "source", "datasource",
          "sourcetype", "txt",
          "event", "Hello world!"
      )
  );

  when(request.getReader()).thenReturn(reader);
  this.task.eventServlet.doPost(request, response);
  verifyResponse(response, 200, Json.MEDIA_TYPE);
  assertEquals(1, this.task.sourceRecordConcurrentLinkedDeque.size(), "Size does not match.");
}
 

/** Tests that {@link BigQueryServicesImpl.JobServiceImpl#startLoadJob} succeeds. */
@Test
public void testStartLoadJobSucceeds() throws IOException, InterruptedException {
  Job testJob = new Job();
  JobReference jobRef = new JobReference();
  jobRef.setJobId("jobId");
  jobRef.setProjectId("projectId");
  testJob.setJobReference(jobRef);

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

  Sleeper sleeper = new FastNanoClockAndSleeper();
  JobServiceImpl.startJob(
      testJob,
      new ApiErrorExtractor(),
      bigquery,
      sleeper,
      BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()));

  verify(response, times(1)).getStatusCode();
  verify(response, times(1)).getContent();
  verify(response, times(1)).getContentType();
  expectedLogs.verifyInfo(String.format("Started BigQuery job: %s", jobRef));
}
 

/** Tests that {@link BigQueryServicesImpl.JobServiceImpl#startLoadJob} succeeds with a retry. */
@Test
public void testStartLoadJobRetry() throws IOException, InterruptedException {
  Job testJob = new Job();
  JobReference jobRef = new JobReference();
  jobRef.setJobId("jobId");
  jobRef.setProjectId("projectId");
  testJob.setJobReference(jobRef);

  // 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(testJob));

  Sleeper sleeper = new FastNanoClockAndSleeper();
  JobServiceImpl.startJob(
      testJob,
      new ApiErrorExtractor(),
      bigquery,
      sleeper,
      BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()));

  verify(response, times(2)).getStatusCode();
  verify(response, times(2)).getContent();
  verify(response, times(2)).getContentType();
}
 

/** Tests that {@link BigQueryServicesImpl.JobServiceImpl#pollJob} succeeds. */
@Test
public void testPollJobSucceeds() throws IOException, InterruptedException {
  Job testJob = new Job();
  testJob.setStatus(new JobStatus().setState("DONE"));

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

  BigQueryServicesImpl.JobServiceImpl jobService =
      new BigQueryServicesImpl.JobServiceImpl(bigquery);
  JobReference jobRef = new JobReference().setProjectId("projectId").setJobId("jobId");
  Job job = jobService.pollJob(jobRef, Sleeper.DEFAULT, BackOff.ZERO_BACKOFF);

  assertEquals(testJob, job);
  verify(response, times(1)).getStatusCode();
  verify(response, times(1)).getContent();
  verify(response, times(1)).getContentType();
}
 

/** Tests that {@link BigQueryServicesImpl.JobServiceImpl#pollJob} fails. */
@Test
public void testPollJobFailed() throws IOException, InterruptedException {
  Job testJob = new Job();
  testJob.setStatus(new JobStatus().setState("DONE").setErrorResult(new ErrorProto()));

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

  BigQueryServicesImpl.JobServiceImpl jobService =
      new BigQueryServicesImpl.JobServiceImpl(bigquery);
  JobReference jobRef = new JobReference().setProjectId("projectId").setJobId("jobId");
  Job job = jobService.pollJob(jobRef, Sleeper.DEFAULT, BackOff.ZERO_BACKOFF);

  assertEquals(testJob, job);
  verify(response, times(1)).getStatusCode();
  verify(response, times(1)).getContent();
  verify(response, times(1)).getContentType();
}
 

/**
 * Convert the passed Json "result" to a response.
 *
 * <p>A {@code null} result will be converted to a successful response. An error response will
 * be generated only if the {@code result} is a {@link GoogleJsonError}.
 *
 * @param result the Json execute result.
 * @return the converted response of the result.
 * @throws IOException
 */
private MockLowLevelHttpResponse getResponse(GenericJson result) throws IOException {
  MockLowLevelHttpResponse response = new MockLowLevelHttpResponse();
  if (result instanceof GoogleJsonError) {
    GoogleJsonError error = (GoogleJsonError) result;
    String errorContent = JSON_FACTORY.toString(new GenericJson().set("error", error));
    response
        .setStatusCode(error.getCode())
        .setReasonPhrase(error.getMessage())
        .setContentType(Json.MEDIA_TYPE)
        .setContent(errorContent);
  } else {
    response
        .setStatusCode(HTTP_OK)
        .setContentType(Json.MEDIA_TYPE)
        .setContent(result == null ? "" : JSON_FACTORY.toString(result));
  }
  return response;
}
 

public void testExecute_void() throws Exception {
  HttpTransport transport = new MockHttpTransport() {
    @Override
    public LowLevelHttpRequest buildRequest(final String method, final String url) {
      return new MockLowLevelHttpRequest() {
        @Override
        public LowLevelHttpResponse execute() {
          return new MockLowLevelHttpResponse().setContent("{\"a\":\"ignored\"}")
              .setContentType(Json.MEDIA_TYPE);
        }
      };
    }
  };
  MockGoogleClient client = new MockGoogleClient.Builder(
      transport, ROOT_URL, SERVICE_PATH, JSON_OBJECT_PARSER, null).setApplicationName(
      "Test Application").build();
  MockGoogleClientRequest<Void> request =
      new MockGoogleClientRequest<Void>(client, HttpMethods.GET, URI_TEMPLATE, null, Void.class);
  Void v = request.execute();
  assertNull(v);
}
 

@Test
public void testGetTableNotFound() throws IOException, InterruptedException {
  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(404);

  BigQueryServicesImpl.DatasetServiceImpl datasetService =
      new BigQueryServicesImpl.DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());

  TableReference tableRef =
      new TableReference()
          .setProjectId("projectId")
          .setDatasetId("datasetId")
          .setTableId("tableId");
  Table table = datasetService.getTable(tableRef, null, BackOff.ZERO_BACKOFF, Sleeper.DEFAULT);

  assertNull(table);
  verify(response, times(1)).getStatusCode();
  verify(response, times(1)).getContent();
  verify(response, times(1)).getContentType();
}
 

@Test
public void testIsTableEmptySucceeds() throws Exception {
  TableReference tableRef =
      new TableReference()
          .setProjectId("projectId")
          .setDatasetId("datasetId")
          .setTableId("tableId");

  TableDataList testDataList = new TableDataList().setRows(ImmutableList.of(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(testDataList));

  BigQueryServicesImpl.DatasetServiceImpl datasetService =
      new BigQueryServicesImpl.DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());

  assertFalse(datasetService.isTableEmpty(tableRef, BackOff.ZERO_BACKOFF, Sleeper.DEFAULT));

  verify(response, times(2)).getStatusCode();
  verify(response, times(2)).getContent();
  verify(response, times(2)).getContentType();
}
 

/** Tests that {@link BigQueryServicesImpl.JobServiceImpl#pollJob} returns UNKNOWN. */
@Test
public void testPollJobUnknown() throws IOException, InterruptedException {
  Job testJob = new Job();
  testJob.setStatus(new JobStatus());

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

  BigQueryServicesImpl.JobServiceImpl jobService =
      new BigQueryServicesImpl.JobServiceImpl(bigquery);
  JobReference jobRef = new JobReference().setProjectId("projectId").setJobId("jobId");
  Job job = jobService.pollJob(jobRef, Sleeper.DEFAULT, BackOff.STOP_BACKOFF);

  assertEquals(null, job);
  verify(response, times(1)).getStatusCode();
  verify(response, times(1)).getContent();
  verify(response, times(1)).getContentType();
}
 

@Override
public LowLevelHttpRequest buildRequest(String name, String url) {
  return new MockLowLevelHttpRequest() {
      @Override
    public LowLevelHttpResponse execute() {
      MockLowLevelHttpResponse r = new MockLowLevelHttpResponse();
      r.setStatusCode(200);
      r.addHeader("Cache-Control", "max-age=" + MAX_AGE);
      if (useAgeHeader) {
        r.addHeader("Age", String.valueOf(AGE));
      }
      r.setContentType(Json.MEDIA_TYPE);
      r.setContent(TEST_CERTIFICATES);
      return r;
    }
  };
}
 

public static MockLowLevelHttpResponse jsonErrorResponse(ErrorResponses errorResponse)
    throws IOException {
  GoogleJsonError.ErrorInfo errorInfo = new GoogleJsonError.ErrorInfo();
  errorInfo.setReason(errorResponse.getErrorReason());
  errorInfo.setDomain(errorResponse.getErrorDomain());
  errorInfo.setFactory(JSON_FACTORY);

  GoogleJsonError jsonError = new GoogleJsonError();
  jsonError.setCode(errorResponse.getErrorCode());
  jsonError.setErrors(ImmutableList.of(errorInfo));
  jsonError.setMessage(errorResponse.getErrorMessage());
  jsonError.setFactory(JSON_FACTORY);

  GenericJson errorResponseJson = new GenericJson();
  errorResponseJson.set("error", jsonError);
  errorResponseJson.setFactory(JSON_FACTORY);

  return new MockLowLevelHttpResponse()
      .setContent(errorResponseJson.toPrettyString())
      .setContentType(Json.MEDIA_TYPE)
      .setStatusCode(errorResponse.getResponseCode());
}
 

public void testDownload() throws Exception {
  HttpTransport transport =
      new MockHttpTransport() {
        @Override
        public LowLevelHttpRequest buildRequest(String method, String url) throws IOException {
          return new MockLowLevelHttpRequest() {
            @Override
            public LowLevelHttpResponse execute() throws IOException {
              MockLowLevelHttpResponse result = new MockLowLevelHttpResponse();
              result.setContentType(Json.MEDIA_TYPE);
              result.setContent(SAMPLE);
              return result;
            }
          };
        }
      };
  HttpRequest request =
      transport.createRequestFactory().buildGetRequest(HttpTesting.SIMPLE_GENERIC_URL);
  HttpResponse response = request.execute();
  ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
  response.download(outputStream);
  assertEquals(SAMPLE, outputStream.toString("UTF-8"));
}
 

public void testParseAsString_utf8() throws Exception {
  HttpTransport transport =
      new MockHttpTransport() {
        @Override
        public LowLevelHttpRequest buildRequest(String method, String url) throws IOException {
          return new MockLowLevelHttpRequest() {
            @Override
            public LowLevelHttpResponse execute() throws IOException {
              MockLowLevelHttpResponse result = new MockLowLevelHttpResponse();
              result.setContentType(Json.MEDIA_TYPE);
              result.setContent(SAMPLE);
              return result;
            }
          };
        }
      };
  HttpRequest request =
      transport.createRequestFactory().buildGetRequest(HttpTesting.SIMPLE_GENERIC_URL);
  HttpResponse response = request.execute();
  assertEquals(SAMPLE, response.parseAsString());
}
 

/** Returns a valid GoogleJsonResponseException for the given status code and error message.  */
private GoogleJsonResponseException makeResponseException(
    final int statusCode,
    final String message) throws Exception {
  HttpTransport transport = new MockHttpTransport() {
    @Override
    public LowLevelHttpRequest buildRequest(String method, String url) {
      return new MockLowLevelHttpRequest() {
        @Override
        public LowLevelHttpResponse execute() {
          MockLowLevelHttpResponse response = new MockLowLevelHttpResponse();
          response.setStatusCode(statusCode);
          response.setContentType(Json.MEDIA_TYPE);
          response.setContent(String.format(
              "{\"error\":{\"code\":%d,\"message\":\"%s\",\"domain\":\"global\","
              + "\"reason\":\"duplicate\"}}",
              statusCode,
              message));
          return response;
        }};
    }};
  HttpRequest request = transport.createRequestFactory()
      .buildGetRequest(HttpTesting.SIMPLE_GENERIC_URL)
      .setThrowExceptionOnExecuteError(false);
  return GoogleJsonResponseException.from(new JacksonFactory(), request.execute());
}
 

@Test
public void testCreateTableSucceeds() throws IOException {
  TableReference ref =
      new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
  Table testTable = new Table().setTableReference(ref);
  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(200);
  when(response.getContent()).thenReturn(toStream(testTable));

  BigQueryServicesImpl.DatasetServiceImpl services =
      new BigQueryServicesImpl.DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());
  Table ret =
      services.tryCreateTable(
          testTable, new RetryBoundedBackOff(0, BackOff.ZERO_BACKOFF), Sleeper.DEFAULT);
  assertEquals(testTable, ret);
  verify(response, times(1)).getStatusCode();
  verify(response, times(1)).getContent();
  verify(response, times(1)).getContentType();
}
 

@Test
public void testExecuteWithRetries() throws IOException, InterruptedException {
  Table testTable = new Table();

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

  Table table =
      BigQueryServicesImpl.executeWithRetries(
          bigquery.tables().get("projectId", "datasetId", "tableId"),
          "Failed to get table.",
          Sleeper.DEFAULT,
          BackOff.STOP_BACKOFF,
          BigQueryServicesImpl.ALWAYS_RETRY);

  assertEquals(testTable, table);
  verify(response, times(1)).getStatusCode();
  verify(response, times(1)).getContent();
  verify(response, times(1)).getContentType();
}
 

@Test
public void testIsTableEmptyThrows() throws Exception {
  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(401);

  TableReference tableRef =
      new TableReference()
          .setProjectId("projectId")
          .setDatasetId("datasetId")
          .setTableId("tableId");

  BigQueryServicesImpl.DatasetServiceImpl datasetService =
      new BigQueryServicesImpl.DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());

  thrown.expect(IOException.class);
  thrown.expectMessage(String.format("Unable to list table data: %s", tableRef.getTableId()));

  datasetService.isTableEmpty(tableRef, BackOff.STOP_BACKOFF, Sleeper.DEFAULT);
}
 

@Test
public void testIsTableEmptyNoRetryForNotFound() throws IOException, InterruptedException {
  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(404);

  BigQueryServicesImpl.DatasetServiceImpl datasetService =
      new BigQueryServicesImpl.DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());

  TableReference tableRef =
      new TableReference()
          .setProjectId("projectId")
          .setDatasetId("datasetId")
          .setTableId("tableId");

  thrown.expect(IOException.class);
  thrown.expectMessage(String.format("Unable to list table data: %s", tableRef.getTableId()));

  try {
    datasetService.isTableEmpty(tableRef, BackOff.ZERO_BACKOFF, Sleeper.DEFAULT);
  } finally {
    verify(response, times(1)).getStatusCode();
    verify(response, times(1)).getContent();
    verify(response, times(1)).getContentType();
  }
}
 

@Test
public void testGetTableThrows() throws Exception {
  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(401);

  TableReference tableRef =
      new TableReference()
          .setProjectId("projectId")
          .setDatasetId("datasetId")
          .setTableId("tableId");

  thrown.expect(IOException.class);
  thrown.expectMessage(String.format("Unable to get table: %s", tableRef.getTableId()));

  BigQueryServicesImpl.DatasetServiceImpl datasetService =
      new BigQueryServicesImpl.DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());
  datasetService.getTable(tableRef, null, BackOff.STOP_BACKOFF, Sleeper.DEFAULT);
}
 

/** Tests that {@link BigQueryServicesImpl} retries quota rate limited attempts. */
@Test
public void testCreateTableRetry() throws IOException {
  TableReference ref =
      new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
  Table testTable = new Table().setTableReference(ref);

  // 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(testTable));

  BigQueryServicesImpl.DatasetServiceImpl services =
      new BigQueryServicesImpl.DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());
  Table ret =
      services.tryCreateTable(
          testTable, new RetryBoundedBackOff(3, BackOff.ZERO_BACKOFF), Sleeper.DEFAULT);
  assertEquals(testTable, ret);
  verify(response, times(2)).getStatusCode();
  verify(response, times(2)).getContent();
  verify(response, times(2)).getContentType();
  verifyNotNull(ret.getTableReference());
  expectedLogs.verifyInfo(
      "Quota limit reached when creating table project:dataset.table, "
          + "retrying up to 5.0 minutes");
}
 

public HttpResponse createUserForHttpResponse(java.io.InputStream body, String mediaType) throws IOException {
    // verify the required parameter 'body' is set
        if (body == null) {
        throw new IllegalArgumentException("Missing the required parameter 'body' when calling createUser");
        }
        UriBuilder uriBuilder = UriBuilder.fromUri(apiClient.getBasePath() + "/user");

        String localVarUrl = uriBuilder.build().toString();
        GenericUrl genericUrl = new GenericUrl(localVarUrl);

        HttpContent content = body == null ?
          apiClient.new JacksonJsonHttpContent(null) :
          new InputStreamContent(mediaType == null ? Json.MEDIA_TYPE : mediaType, body);
        return apiClient.getHttpRequestFactory().buildRequest(HttpMethods.POST, genericUrl, content).execute();
}
 

/** Tests that {@link BigQueryServicesImpl} does not retry non-rate-limited attempts. */
@Test
public void testCreateTableDoesNotRetry() throws IOException {
  TableReference ref =
      new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
  Table testTable = new Table().setTableReference(ref);
  // First response is 403 not-rate-limited, 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(testTable));

  thrown.expect(GoogleJsonResponseException.class);
  thrown.expectMessage("actually forbidden");

  BigQueryServicesImpl.DatasetServiceImpl services =
      new BigQueryServicesImpl.DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());
  try {
    services.tryCreateTable(
        testTable, new RetryBoundedBackOff(3, BackOff.ZERO_BACKOFF), Sleeper.DEFAULT);
    fail();
  } catch (IOException e) {
    verify(response, times(1)).getStatusCode();
    verify(response, times(1)).getContent();
    verify(response, times(1)).getContentType();
    throw e;
  }
}
 

/**
 * 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:");
}
 

/** Builds a fake GoogleJsonResponseException for testing API error handling. */
private static GoogleJsonResponseException googleJsonResponseException(
    final int status, final String reason, final String message) throws IOException {
  final JsonFactory jsonFactory = new JacksonFactory();
  HttpTransport transport =
      new MockHttpTransport() {
        @Override
        public LowLevelHttpRequest buildRequest(String method, String url) throws IOException {
          ErrorInfo errorInfo = new ErrorInfo();
          errorInfo.setReason(reason);
          errorInfo.setMessage(message);
          errorInfo.setFactory(jsonFactory);
          GenericJson error = new GenericJson();
          error.set("code", status);
          error.set("errors", Arrays.asList(errorInfo));
          error.setFactory(jsonFactory);
          GenericJson errorResponse = new GenericJson();
          errorResponse.set("error", error);
          errorResponse.setFactory(jsonFactory);
          return new MockLowLevelHttpRequest()
              .setResponse(
                  new MockLowLevelHttpResponse()
                      .setContent(errorResponse.toPrettyString())
                      .setContentType(Json.MEDIA_TYPE)
                      .setStatusCode(status));
        }
      };
  HttpRequest request =
      transport.createRequestFactory().buildGetRequest(HttpTesting.SIMPLE_GENERIC_URL);
  request.setThrowExceptionOnExecuteError(false);
  HttpResponse response = request.execute();
  return GoogleJsonResponseException.from(jsonFactory, response);
}
 

private LowLevelHttpResponse generateMockResponse(WorkItem... workItems) throws Exception {
  MockLowLevelHttpResponse response = new MockLowLevelHttpResponse();
  response.setContentType(Json.MEDIA_TYPE);
  LeaseWorkItemResponse lease = new LeaseWorkItemResponse();
  lease.setWorkItems(Lists.newArrayList(workItems));
  // N.B. Setting the factory is necessary in order to get valid JSON.
  lease.setFactory(Transport.getJsonFactory());
  response.setContent(lease.toPrettyString());
  return response;
}