Examples with SplitReadStreamRequest com.google.cloud.bigquery.storage.v1.SplitReadStreamRequest used on opensource projects

Example 1 with SplitReadStreamRequest

use of com.google.cloud.bigquery.storage.v1.SplitReadStreamRequest in project beam by apache.

the class BigQueryServicesImplTest method testSplitReadStreamSetsRequestCountMetricOnError.

@Test
public void testSplitReadStreamSetsRequestCountMetricOnError() throws InterruptedException, IOException {
    BigQueryServices.StorageClient client = mock(BigQueryServicesImpl.StorageClientImpl.class);
    SplitReadStreamRequest request = null;
    StatusCode statusCode = new StatusCode() {

        @Override
        public Code getCode() {
            return Code.RESOURCE_EXHAUSTED;
        }

        @Override
        public Object getTransportCode() {
            return null;
        }
    };
    when(client.splitReadStream(request)).thenThrow(new ApiException("Resource Exhausted", null, statusCode, // Mock implementation.
    false));
    // Real implementation.
    when(client.splitReadStream(any(), any())).thenCallRealMethod();
    thrown.expect(ApiException.class);
    thrown.expectMessage("Resource Exhausted");
    client.splitReadStream(request, "myproject:mydataset.mytable");
    verifyReadMetricWasSet("myproject", "mydataset", "mytable", "resource_exhausted", 1);
}

Example 2 with SplitReadStreamRequest

use of com.google.cloud.bigquery.storage.v1.SplitReadStreamRequest in project beam by apache.

the class BigQueryServicesImplTest method testSplitReadStreamSetsRequestCountMetric.

@Test
public void testSplitReadStreamSetsRequestCountMetric() throws InterruptedException, IOException {
    BigQueryServices.StorageClient client = mock(BigQueryServicesImpl.StorageClientImpl.class);
    SplitReadStreamRequest request = null;
    when(client.splitReadStream(request)).thenReturn(// Mock implementation.
    SplitReadStreamResponse.newBuilder().build());
    // Real implementation.
    when(client.splitReadStream(any(), any())).thenCallRealMethod();
    client.splitReadStream(request, "myproject:mydataset.mytable");
    verifyReadMetricWasSet("myproject", "mydataset", "mytable", "ok", 1);
}