use of com.google.api.client.testing.util.MockSleeper in project elasticsearch by elastic.
the class RetryHttpInitializerWrapperTests method testIOExceptionRetry.
public void testIOExceptionRetry() throws Exception {
FailThenSuccessBackoffTransport fakeTransport = new FailThenSuccessBackoffTransport(HttpStatusCodes.STATUS_CODE_SERVER_ERROR, 1, true);
MockGoogleCredential credential = RetryHttpInitializerWrapper.newMockCredentialBuilder().build();
MockSleeper mockSleeper = new MockSleeper();
RetryHttpInitializerWrapper retryHttpInitializerWrapper = new RetryHttpInitializerWrapper(credential, mockSleeper, TimeValue.timeValueMillis(500));
Compute client = new Compute.Builder(fakeTransport, new JacksonFactory(), null).setHttpRequestInitializer(retryHttpInitializerWrapper).setApplicationName("test").build();
HttpRequest request = client.getRequestFactory().buildRequest("Get", new GenericUrl("http://elasticsearch.com"), null);
HttpResponse response = request.execute();
assertThat(mockSleeper.getCount(), equalTo(1));
assertThat(response.getStatusCode(), equalTo(200));
}
use of com.google.api.client.testing.util.MockSleeper in project beam by apache.
the class BigQueryServicesImplTest method testInsertRetry.
/**
* Tests that {@link DatasetServiceImpl#insertAll} retries quota rate limited attempts.
*/
@Test
public void testInsertRetry() throws Exception {
TableReference ref = new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
List<ValueInSingleWindow<TableRow>> rows = new ArrayList<>();
rows.add(wrapTableRow(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);
verify(response, times(2)).getStatusCode();
verify(response, times(2)).getContent();
verify(response, times(2)).getContentType();
expectedLogs.verifyInfo("BigQuery insertAll exceeded rate limit, retrying");
}
use of com.google.api.client.testing.util.MockSleeper in project beam by apache.
the class BigQueryServicesImplTest method testInsertRetrySelectRows.
/**
* 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<FailsafeValueInSingleWindow<TableRow, 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();
setupMockResponses(response -> {
when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
when(response.getStatusCode()).thenReturn(200);
when(response.getContent()).thenReturn(toStream(bFailed));
}, response -> {
when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
when(response.getStatusCode()).thenReturn(200);
when(response.getContent()).thenReturn(toStream(allRowsSucceeded));
});
DatasetServiceImpl dataService = new DatasetServiceImpl(bigquery, null, PipelineOptionsFactory.create());
dataService.insertAll(ref, rows, insertIds, BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()), TEST_BACKOFF, new MockSleeper(), InsertRetryPolicy.alwaysRetry(), null, null, false, false, false, null);
verifyAllResponsesAreRead();
verifyWriteMetricWasSet("project", "dataset", "table", "unknown", 1);
verifyWriteMetricWasSet("project", "dataset", "table", "ok", 1);
}
use of com.google.api.client.testing.util.MockSleeper in project beam by apache.
the class BigQueryServicesImplTest method testSkipInvalidRowsIgnoreUnknownIgnoreInsertIdsValuesStreaming.
/**
* 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<FailsafeValueInSingleWindow<TableRow, TableRow>> rows = ImmutableList.of(wrapValue(new TableRow()), wrapValue(new TableRow()));
final TableDataInsertAllResponse allRowsSucceeded = new TableDataInsertAllResponse();
// Return a 200 response each time
MockSetupFunction allRowsSucceededResponseFunction = response -> {
when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
when(response.getStatusCode()).thenReturn(200);
when(response.getContent()).thenReturn(toStream(allRowsSucceeded));
};
setupMockResponses(allRowsSucceededResponseFunction, allRowsSucceededResponseFunction);
DatasetServiceImpl dataService = new DatasetServiceImpl(bigquery, null, PipelineOptionsFactory.create());
// First, test with all flags disabled
dataService.insertAll(ref, rows, null, BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()), TEST_BACKOFF, new MockSleeper(), InsertRetryPolicy.neverRetry(), Lists.newArrayList(), ErrorContainer.TABLE_ROW_ERROR_CONTAINER, false, false, false, null);
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()), TEST_BACKOFF, new MockSleeper(), InsertRetryPolicy.neverRetry(), Lists.newArrayList(), ErrorContainer.TABLE_ROW_ERROR_CONTAINER, true, true, true, null);
parsedRequest = fromString(request.getContentAsString(), TableDataInsertAllRequest.class);
assertTrue(parsedRequest.getSkipInvalidRows());
assertTrue(parsedRequest.getIgnoreUnknownValues());
assertNull(parsedRequest.getRows().get(0).getInsertId());
assertNull(parsedRequest.getRows().get(1).getInsertId());
verifyWriteMetricWasSet("project", "dataset", "table", "ok", 2);
}
use of com.google.api.client.testing.util.MockSleeper in project beam by apache.
the class BigQueryServicesImplTest method testExtendedErrorRetrieval.
/**
* 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<FailsafeValueInSingleWindow<TableRow, 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(wrapErrorValue(new BigQueryInsertError(rows.get(0).getValue(), failures.getInsertErrors().get(0), ref)), wrapErrorValue(new BigQueryInsertError(rows.get(1).getValue(), failures.getInsertErrors().get(1), ref)));
setupMockResponses(response -> {
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, null, PipelineOptionsFactory.create());
List<ValueInSingleWindow<BigQueryInsertError>> failedInserts = Lists.newArrayList();
dataService.insertAll(ref, rows, null, BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()), TEST_BACKOFF, new MockSleeper(), InsertRetryPolicy.neverRetry(), failedInserts, ErrorContainer.BIG_QUERY_INSERT_ERROR_ERROR_CONTAINER, false, false, false, null);
assertThat(failedInserts, is(expected));
}
Aggregations