fix: add test case

googleapis · May 8, 2020 · 360ef1e · 360ef1e
1 parent 55f1a05
commit 360ef1e
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 8 deletions.
diff --git a/...main/java/com/google/cloud/bigquery/storage/v1/stub/readrows/ApiResultRetryAlgorithm.java b/...main/java/com/google/cloud/bigquery/storage/v1/stub/readrows/ApiResultRetryAlgorithm.java
@@ -20,7 +20,6 @@
 import com.google.api.gax.retrying.ResultRetryAlgorithm;
 import com.google.api.gax.retrying.TimedAttemptSettings;
 import com.google.api.gax.rpc.ApiException;
-import com.google.api.gax.rpc.DeadlineExceededException;
 import io.grpc.Status;
 import org.threeten.bp.Duration;
 
@@ -33,7 +32,7 @@ public class ApiResultRetryAlgorithm<ResponseT> implements ResultRetryAlgorithm<
   @Override
   public TimedAttemptSettings createNextAttempt(
       Throwable prevThrowable, ResponseT prevResponse, TimedAttemptSettings prevSettings) {
-    if (prevThrowable != null && prevThrowable instanceof DeadlineExceededException) {
+    if (prevThrowable != null && isRetryable(prevThrowable)) {
       return TimedAttemptSettings.newBuilder()
           .setGlobalSettings(prevSettings.getGlobalSettings())
           .setRetryDelay(prevSettings.getRetryDelay())
@@ -49,13 +48,18 @@ public TimedAttemptSettings createNextAttempt(
   @Override
   public boolean shouldRetry(Throwable prevThrowable, ResponseT prevResponse) {
     if (prevThrowable != null) {
-      Status status = Status.fromThrowable(prevThrowable);
-      if (status.getCode() == Status.Code.INTERNAL
-          && status.getDescription() != null
-          && status.getDescription().equals("Received unexpected EOS on DATA frame from server")) {
-        return true;
-      }
+      return isRetryable(prevThrowable);
     }
     return (prevThrowable instanceof ApiException) && ((ApiException) prevThrowable).isRetryable();
   }
+
+  private boolean isRetryable(Throwable prevThrowable) {
+    Status status = Status.fromThrowable(prevThrowable);
+    if (status.getCode() == Status.Code.INTERNAL
+        && status.getDescription() != null
+        && status.getDescription().equals("Received unexpected EOS on DATA frame from server")) {
+      return true;
+    }
+    return false;
+  }
 }
diff --git a/...erystorage/src/test/java/com/google/cloud/bigquery/storage/v1/BigQueryReadClientTest.java b/...erystorage/src/test/java/com/google/cloud/bigquery/storage/v1/BigQueryReadClientTest.java
@@ -162,4 +162,24 @@ public void readRowsExceptionTest() throws Exception {
       Assert.assertEquals(StatusCode.Code.INVALID_ARGUMENT, apiException.getStatusCode().getCode());
     }
   }
+
+  @Test
+  @SuppressWarnings("all")
+  public void readRowsRetryingExceptionTest() throws ExecutionException, InterruptedException {
+    StatusRuntimeException exception =
+        new StatusRuntimeException(
+            Status.INTERNAL.withDescription("Received unexpected EOS on DATA frame from server"));
+    mockBigQueryRead.addException(exception);
+    long rowCount = 1340416618L;
+    ReadRowsResponse expectedResponse = ReadRowsResponse.newBuilder().setRowCount(rowCount).build();
+    mockBigQueryRead.addResponse(expectedResponse);
+    ReadRowsRequest request = ReadRowsRequest.newBuilder().build();
+
+    MockStreamObserver<ReadRowsResponse> responseObserver = new MockStreamObserver<>();
+
+    ServerStreamingCallable<ReadRowsRequest, ReadRowsResponse> callable = client.readRowsCallable();
+    callable.serverStreamingCall(request, responseObserver);
+    List<ReadRowsResponse> actualResponses = responseObserver.future().get();
+    Assert.assertEquals(1, actualResponses.size());
+  }
 }