fix: set timeouts for BatchGetDocuments/RunQuery (#799)

schmidt-sebastian · web-flow · commit 6cd2a4526561 · 2021-10-27T12:55:21.000-06:00
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
@@ -132,6 +132,9 @@ public GrpcFirestoreRpc(final FirestoreOptions options) throws IOException {
       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/ITBulkWriterTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITBulkWriterTest.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.google.cloud.firestore;
+package com.google.cloud.firestore.it;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
@@ -24,6 +24,14 @@
 
 import com.google.api.core.ApiFuture;
 import com.google.cloud.Timestamp;
+import com.google.cloud.firestore.BulkWriter;
+import com.google.cloud.firestore.CollectionReference;
+import com.google.cloud.firestore.DocumentReference;
+import com.google.cloud.firestore.DocumentSnapshot;
+import com.google.cloud.firestore.Firestore;
+import com.google.cloud.firestore.FirestoreOptions;
+import com.google.cloud.firestore.LocalFirestoreHelper;
+import com.google.cloud.firestore.WriteResult;
 import com.google.common.base.Preconditions;
 import java.util.ArrayList;
 import java.util.Collections;
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
@@ -58,7 +58,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
@@ -32,7 +32,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;
@@ -100,7 +102,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;
@@ -1761,4 +1767,41 @@ public void testRecursiveDeleteWithCustomBulkWriterInstance() throws Exception {
     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());
+  }
 }
