Examples with MockSleeper com.google.api.client.testing.util.MockSleeper used on opensource projects

Example 11 with MockSleeper

use of com.google.api.client.testing.util.MockSleeper in project beam by apache.

the class BigQueryServicesImplTest method testInsertWithinRowCountLimits.

/**
 * Tests that {@link DatasetServiceImpl#insertAll} does not go over limit of rows per request.
 */
@Test
public void testInsertWithinRowCountLimits() 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")), wrapValue(new TableRow().set("row", "c")));
    List<String> insertIds = ImmutableList.of("a", "b", "c");
    final TableDataInsertAllResponse allRowsSucceeded = new TableDataInsertAllResponse();
    setupMockResponses(response -> {
        when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
        when(response.getStatusCode()).thenReturn(200);
        when(response.getContent()).thenReturn(toStream(allRowsSucceeded));
    }, response -> {
        when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
        when(response.getStatusCode()).thenReturn(200);
        when(response.getContent()).thenReturn(toStream(allRowsSucceeded));
    }, 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.fromArgs("--maxStreamingRowsToBatch=1").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", "ok", 3);
}

Example 12 with MockSleeper

use of com.google.api.client.testing.util.MockSleeper in project elasticsearch by elastic.

the class RetryHttpInitializerWrapperTests method testRetryWaitTooLong.

public void testRetryWaitTooLong() throws Exception {
    TimeValue maxWaitTime = TimeValue.timeValueMillis(10);
    int maxRetryTimes = 50;
    FailThenSuccessBackoffTransport fakeTransport = new FailThenSuccessBackoffTransport(HttpStatusCodes.STATUS_CODE_SERVER_ERROR, maxRetryTimes);
    JsonFactory jsonFactory = new JacksonFactory();
    MockGoogleCredential credential = RetryHttpInitializerWrapper.newMockCredentialBuilder().build();
    MockSleeper oneTimeSleeper = new MockSleeper() {

        @Override
        public void sleep(long millis) throws InterruptedException {
            Thread.sleep(maxWaitTime.getMillis());
            // important number, use this to get count
            super.sleep(0);
        }
    };
    RetryHttpInitializerWrapper retryHttpInitializerWrapper = new RetryHttpInitializerWrapper(credential, oneTimeSleeper, maxWaitTime);
    Compute client = new Compute.Builder(fakeTransport, jsonFactory, null).setHttpRequestInitializer(retryHttpInitializerWrapper).setApplicationName("test").build();
    HttpRequest request1 = client.getRequestFactory().buildRequest("Get", new GenericUrl("http://elasticsearch.com"), null);
    try {
        request1.execute();
        fail("Request should fail if wait too long");
    } catch (HttpResponseException e) {
        assertThat(e.getStatusCode(), equalTo(HttpStatusCodes.STATUS_CODE_SERVER_ERROR));
        // should only retry once.
        assertThat(oneTimeSleeper.getCount(), lessThan(maxRetryTimes));
    }
}

Example 13 with MockSleeper

use of com.google.api.client.testing.util.MockSleeper in project elasticsearch by elastic.

the class RetryHttpInitializerWrapperTests method testSimpleRetry.

public void testSimpleRetry() throws Exception {
    FailThenSuccessBackoffTransport fakeTransport = new FailThenSuccessBackoffTransport(HttpStatusCodes.STATUS_CODE_SERVER_ERROR, 3);
    MockGoogleCredential credential = RetryHttpInitializerWrapper.newMockCredentialBuilder().build();
    MockSleeper mockSleeper = new MockSleeper();
    RetryHttpInitializerWrapper retryHttpInitializerWrapper = new RetryHttpInitializerWrapper(credential, mockSleeper, TimeValue.timeValueSeconds(5));
    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(3));
    assertThat(response.getStatusCode(), equalTo(200));
}

Example 14 with MockSleeper

use of com.google.api.client.testing.util.MockSleeper in project beam by apache.

the class BigQueryServicesImplTest method testInsertDoesNotRetry.

/**
   * Tests that {@link DatasetServiceImpl#insertAll} does not retry non-rate-limited attempts.
   */
@Test
public void testInsertDoesNotRetry() throws Throwable {
    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 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(new TableDataInsertAllResponse()));
    thrown.expect(GoogleJsonResponseException.class);
    thrown.expectMessage("actually forbidden");
    DatasetServiceImpl dataService = new DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());
    try {
        dataService.insertAll(ref, rows, null, BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()), new MockSleeper(), InsertRetryPolicy.alwaysRetry(), null);
        fail();
    } catch (RuntimeException e) {
        verify(response, times(1)).getStatusCode();
        verify(response, times(1)).getContent();
        verify(response, times(1)).getContentType();
        throw e.getCause();
    }
}

Example 15 with MockSleeper

use of com.google.api.client.testing.util.MockSleeper in project beam by apache.

the class BigQueryServicesImplTest method testSimpleErrorRetrieval.

/**
 * 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<FailsafeValueInSingleWindow<TableRow, TableRow>> rows = ImmutableList.of(wrapValue(new TableRow().set("a", 1)), wrapValue(new TableRow().set("b", 2)));
    final List<ValueInSingleWindow<TableRow>> expected = ImmutableList.of(wrapErrorValue(new TableRow().set("a", 1)), wrapErrorValue(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")))));
    setupMockResponses(response -> {
        when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
        when(response.getStatusCode()).thenReturn(200);
        when(response.getContent()).thenReturn(toStream(failures));
    });
    DatasetServiceImpl dataService = new DatasetServiceImpl(bigquery, null, PipelineOptionsFactory.create());
    List<ValueInSingleWindow<TableRow>> failedInserts = Lists.newArrayList();
    dataService.insertAll(ref, rows, null, BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()), TEST_BACKOFF, new MockSleeper(), InsertRetryPolicy.neverRetry(), failedInserts, ErrorContainer.TABLE_ROW_ERROR_CONTAINER, false, false, false, null);
    assertThat(failedInserts, is(expected));
}