Examples with FailsafeValueInSingleWindow org.apache.beam.sdk.values.FailsafeValueInSingleWindow used on opensource projects

Example 16 with FailsafeValueInSingleWindow

use of org.apache.beam.sdk.values.FailsafeValueInSingleWindow in project beam by apache.

the class BigQueryServicesImplTest method testInsertQuotaExceededRetry.

/**
 * 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<FailsafeValueInSingleWindow<TableRow, TableRow>> rows = new ArrayList<>();
    rows.add(wrapValue(new TableRow()));
    // First response is 403 quota exceeded, second response has valid payload.
    setupMockResponses(response -> {
        when(response.getStatusCode()).thenReturn(403);
        when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
        when(response.getContent()).thenReturn(toStream(errorWithReasonAndStatus("quotaExceeded", 403)));
    }, response -> {
        when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
        when(response.getStatusCode()).thenReturn(200);
        when(response.getContent()).thenReturn(toStream(new TableDataInsertAllResponse()));
    });
    DatasetServiceImpl dataService = new DatasetServiceImpl(bigquery, null, PipelineOptionsFactory.create());
    dataService.insertAll(ref, rows, null, BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()), TEST_BACKOFF, new MockSleeper(), InsertRetryPolicy.alwaysRetry(), null, null, false, false, false, null);
    verifyAllResponsesAreRead();
    expectedLogs.verifyInfo("BigQuery insertAll error, retrying:");
    verifyWriteMetricWasSet("project", "dataset", "table", "quotaexceeded", 1);
}