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);
}
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);
}
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);
}
Aggregations