Search in sources :

Example 1 with StatusCode

use of com.google.api.gax.rpc.StatusCode 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);
}
Also used : SplitReadStreamRequest(com.google.cloud.bigquery.storage.v1.SplitReadStreamRequest) StatusCode(com.google.api.gax.rpc.StatusCode) ApiException(com.google.api.gax.rpc.ApiException) Test(org.junit.Test)

Example 2 with StatusCode

use of com.google.api.gax.rpc.StatusCode in project beam by apache.

the class BigQueryServicesImplTest method testCreateReadSessionSetsRequestCountMetricOnError.

@Test
public void testCreateReadSessionSetsRequestCountMetricOnError() throws InterruptedException, IOException {
    BigQueryServicesImpl.StorageClientImpl client = mock(BigQueryServicesImpl.StorageClientImpl.class);
    CreateReadSessionRequest.Builder builder = CreateReadSessionRequest.newBuilder();
    builder.getReadSessionBuilder().setTable("myproject:mydataset.mytable");
    CreateReadSessionRequest request = builder.build();
    StatusCode statusCode = new StatusCode() {

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

        @Override
        public Object getTransportCode() {
            return null;
        }
    };
    when(client.callCreateReadSession(request)).thenThrow(// Mock implementation.
    new ApiException("Not Found", null, statusCode, false));
    // Real implementation.
    when(client.createReadSession(any())).thenCallRealMethod();
    thrown.expect(ApiException.class);
    thrown.expectMessage("Not Found");
    client.createReadSession(request);
    verifyReadMetricWasSet("myproject", "mydataset", "mytable", "not_found", 1);
}
Also used : StatusCode(com.google.api.gax.rpc.StatusCode) CreateReadSessionRequest(com.google.cloud.bigquery.storage.v1.CreateReadSessionRequest) ApiException(com.google.api.gax.rpc.ApiException) Test(org.junit.Test)

Example 3 with StatusCode

use of com.google.api.gax.rpc.StatusCode in project beam by apache.

the class BigQueryServicesImplTest method testReadRowsSetsRequestCountMetricOnError.

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

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

        @Override
        public Object getTransportCode() {
            return null;
        }
    };
    when(client.readRows(request)).thenThrow(// Mock implementation.
    new ApiException("Internal", null, statusCode, false));
    // Real implementation.
    when(client.readRows(any(), any())).thenCallRealMethod();
    thrown.expect(ApiException.class);
    thrown.expectMessage("Internal");
    client.readRows(request, "myproject:mydataset.mytable");
    verifyReadMetricWasSet("myproject", "mydataset", "mytable", "internal", 1);
}
Also used : ReadRowsRequest(com.google.cloud.bigquery.storage.v1.ReadRowsRequest) StatusCode(com.google.api.gax.rpc.StatusCode) ApiException(com.google.api.gax.rpc.ApiException) Test(org.junit.Test)

Aggregations

ApiException (com.google.api.gax.rpc.ApiException)3 StatusCode (com.google.api.gax.rpc.StatusCode)3 Test (org.junit.Test)3 CreateReadSessionRequest (com.google.cloud.bigquery.storage.v1.CreateReadSessionRequest)1 ReadRowsRequest (com.google.cloud.bigquery.storage.v1.ReadRowsRequest)1 SplitReadStreamRequest (com.google.cloud.bigquery.storage.v1.SplitReadStreamRequest)1