fix: set timeouts for BatchGetDocuments/RunQuery

googleapis · Oct 26, 2021 · e32a6c7 · e32a6c7
1 parent b49a89c
commit e32a6c7
Show file tree

Hide file tree

Showing 4 changed files with 55 additions and 6 deletions.
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/spi/v1/GrpcFirestoreRpc.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/spi/v1/GrpcFirestoreRpc.java
@@ -126,16 +126,16 @@ public GrpcFirestoreRpc(final FirestoreOptions options) throws IOException {
         clientContext = ClientContext.create(settingsBuilder.build());
       }
       ApiFunction<UnaryCallSettings.Builder<?, ?>, Void> retrySettingsSetter =
-          new ApiFunction<UnaryCallSettings.Builder<?, ?>, Void>() {
-            @Override
-            public Void apply(UnaryCallSettings.Builder<?, ?> builder) {
-              builder.setRetrySettings(options.getRetrySettings());
-              return null;
-            }
+          builder -> {
+            builder.setRetrySettings(options.getRetrySettings());
+            return null;
           };
       FirestoreStubSettings.Builder firestoreBuilder =
           FirestoreStubSettings.newBuilder(clientContext)
               .applyToAllUnaryMethods(retrySettingsSetter);
+      // Manually apply the retry settings to streaming methods
+      firestoreBuilder.runQuerySettings().setRetrySettings(options.getRetrySettings());
+      firestoreBuilder.batchGetDocumentsSettings().setRetrySettings(options.getRetrySettings());
       firestoreStub = GrpcFirestoreStub.create(firestoreBuilder.build());
     } catch (Exception e) {
       throw new IOException(e);

diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITQueryWatchTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITQueryWatchTest.java
@@ -59,7 +59,10 @@
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TestName;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
 
+@RunWith(JUnit4.class)
 public final class ITQueryWatchTest {
 
   private static Firestore firestore;

diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITShutdownTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITShutdownTest.java
@@ -36,7 +36,10 @@
 import org.junit.Test;
 import org.junit.rules.TestName;
 import org.junit.rules.Timeout;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
 
+@RunWith(JUnit4.class)
 public class ITShutdownTest {
   @Rule public final Timeout timeout = new Timeout(5, TimeUnit.SECONDS);
   @Rule public TestName testName = new TestName();

diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITSystemTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITSystemTest.java
@@ -33,12 +33,14 @@
 import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
 import com.google.api.core.ApiFuture;
 import com.google.api.core.ApiFutures;
 import com.google.api.core.SettableApiFuture;
+import com.google.api.gax.retrying.RetrySettings;
 import com.google.api.gax.rpc.ApiStreamObserver;
 import com.google.cloud.Timestamp;
 import com.google.cloud.firestore.BulkWriter;
@@ -103,7 +105,11 @@
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TestName;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.threeten.bp.Duration;
 
+@RunWith(JUnit4.class)
 public class ITSystemTest {
 
   private static final double DOUBLE_EPSILON = 0.000001;
@@ -1811,4 +1817,41 @@ public void onResult(DocumentReference documentReference, WriteResult result) {
     assertEquals(0, countCollectionChildren(randomColl));
     assertEquals(6, callbackCount[0]);
   }
+
+  @Test
+  public void testEnforcesTimeouts() throws Exception {
+    FirestoreOptions firestoreOptions =
+        FirestoreOptions.newBuilder()
+            .setRetrySettings(
+                RetrySettings.newBuilder()
+                    .setMaxRpcTimeout(Duration.ofMillis(1))
+                    .setTotalTimeout(Duration.ofMillis(1))
+                    .setInitialRpcTimeout(Duration.ofMillis(1))
+                    .build())
+            .build();
+    firestore = firestoreOptions.getService();
+    CollectionReference collection = firestore.collection("timeout");
+
+    // RunQuery
+    assertThrows(ExecutionException.class, () -> collection.get().get());
+    // CommitRequest
+    assertThrows(ExecutionException.class, () -> collection.add(map()).get());
+    // BulkCommit
+    assertThrows(
+        ExecutionException.class,
+        () -> {
+          BulkWriter bulkWriter = firestore.bulkWriter();
+          ApiFuture<WriteResult> op = bulkWriter.set(collection.document(), map());
+          bulkWriter.close();
+          op.get();
+        });
+    // BatchGetDocuments
+    assertThrows(ExecutionException.class, () -> collection.document().get().get());
+    // ListDocuments
+    assertThrows(FirestoreException.class, () -> collection.listDocuments().iterator().hasNext());
+    // ListCollections
+    assertThrows(
+        FirestoreException.class,
+        () -> collection.document().listCollections().iterator().hasNext());
+  }
 }