fix: set timeouts for BatchGetDocuments/RunQuery

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

google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITQueryTest.java

@@ -0,0 +1,190 @@
+package com.google.cloud.firestore.it;
+
+import static com.google.cloud.firestore.LocalFirestoreHelper.map;
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.cloud.firestore.CollectionReference;
+import com.google.cloud.firestore.Firestore;
+import com.google.cloud.firestore.FirestoreOptions;
+import com.google.cloud.firestore.LocalFirestoreHelper;
+import com.google.cloud.firestore.Query;
+import com.google.cloud.firestore.QueryDocumentSnapshot;
+import com.google.common.base.Preconditions;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ExecutionException;
+import java.util.stream.Collectors;
+import org.junit.After;
+import org.junit.Before;
+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 class ITQueryTest {
+
+  @Rule public TestName testName = new TestName();
+
+  private Firestore firestore;
+  private CollectionReference indexedCollection;
+
+  @Before
+  public void before() {
+    FirestoreOptions firestoreOptions = FirestoreOptions.newBuilder().build();
+    firestore = firestoreOptions.getService();
+    indexedCollection =
+        firestore.collection(
+            String.format("java-%s-%s", testName.getMethodName(), LocalFirestoreHelper.autoId()));
+  }
+
+  @After
+  public void after() throws Exception {
+    Preconditions.checkNotNull(
+        firestore,
+        "Error instantiating Firestore. Check that the service account credentials were properly set.");
+    firestore.close();
+  }
+
+  @Test
+  public void advancedQueries() throws Exception {
+    List<Map<String, Object>> data =
+        new ArrayList<Map<String, Object>>() {
+          {
+            add(map("int", 1, "array", Arrays.asList(1, "foo"))); // d0
+            add(map("array", Arrays.asList(2, "foo"))); // d1
+            add(map("int", 3, "array", Arrays.asList(3, "foo"))); // d2
+            add(map("float", -0.0, "string", "a")); // d3
+            add(map("float", 0, "string", "ab")); // d4
+            add(map("float", 0.0, "string", "b")); // d5
+            add(map("float", Double.NaN)); // d6
+            add(map("multi", true)); // d7
+            add(map("array", "foo")); // d8
+            add(map("array", Collections.singletonList(1))); // d9
+            add(map("multi", 1)); // d10
+            add(map("multi", "string")); // d11
+            add(map("multi", Collections.emptyList())); // d12
+            add(map("a", 0, "b", 0)); // d13
+            add(map("a", 0, "b", 1)); // d14
+            add(map("a", 1, "b", 0)); // d15
+            add(map("a", 1, "b", 1)); // d16
+            add(map()); // d17
+            add(map("null", null)); // d18
+            add(map("prefix", Arrays.asList(1, 2), "suffix", null)); // d19
+            add(map("prefix", Arrays.asList(1), "suffix", 2)); // d20
+            add(map("map", map())); // d20
+            add(map("map", map("value", true))); // d21
+            add(map("map", map("value", false))); // d22
+          }
+        };
+
+    for (int i = 0; i < data.size(); ++i) {
+      addDoc("d" + i, data.get(i));
+    }
+
+    Query q = firestore.collection("indexed");
+
+    verifyResults(q.orderBy("int"), "d0", "d2");
+    verifyResults(q.whereEqualTo("float", Double.NaN), "d6");
+    verifyResults(q.whereEqualTo("float", -0.0), "d3", "d4", "d5");
+    verifyResults(q.whereEqualTo("float", 0), "d3", "d4", "d5");
+    verifyResults(q.whereEqualTo("float", 0.0), "d3", "d4", "d5");
+    verifyResults(q.whereEqualTo("string", "a"), "d3");
+    verifyResults(q.whereGreaterThan("string", "a"), "d4", "d5");
+    verifyResults(q.whereGreaterThanOrEqualTo("string", "a"), "d3", "d4", "d5");
+    verifyResults(q.whereLessThan("string", "b"), "d3", "d4");
+    verifyResults(q.whereLessThan("string", "c"), "d3", "d4", "d5");
+    verifyResults(q.whereGreaterThan("string", "a").whereLessThan("string", "b"), "d4");
+    verifyResults(q.whereArrayContains("array", "foo"), "d0", "d1", "d2");
+    verifyResults(
+        q.whereArrayContainsAny("array", Arrays.asList(1, "foo")), "d0", "d1", "d2", "d9");
+    verifyResults(q.whereGreaterThanOrEqualTo("multi", true), "d7");
+    verifyResults(q.whereGreaterThanOrEqualTo("multi", 0), "d10");
+    verifyResults(q.whereGreaterThanOrEqualTo("multi", ""), "d11");
+    verifyResults(q.whereGreaterThanOrEqualTo("multi", Collections.emptyList()), "d12");
+    verifyResults(q.orderBy("array").startAt(Collections.singletonList(2)), "d1", "d2");
+    verifyResults(
+        q.orderBy("array", Query.Direction.DESCENDING).startAt(Collections.singletonList(2)),
+        "d0",
+        "d8",
+        "d9");
+    verifyResults(
+        q.orderBy("array", Query.Direction.DESCENDING)
+            .startAt(Collections.singletonList(2))
+            .limit(2),
+        "d0",
+        "d9");
+    verifyResults(q.orderBy("array").startAfter(Collections.singletonList(2)), "d1", "d2");
+    verifyResults(
+        q.orderBy("array", Query.Direction.DESCENDING).startAfter(Collections.singletonList(2)),
+        "d0",
+        "d8",
+        "d9");
+    verifyResults(
+        q.orderBy("array", Query.Direction.DESCENDING)
+            .startAfter(Collections.singletonList(2))
+            .limit(2),
+        "d0",
+        "d9");
+    verifyResults(q.orderBy("array").startAfter(Arrays.asList(2, "foo")), "d2");
+    verifyResults(
+        q.orderBy("array", Query.Direction.DESCENDING).startAfter(Arrays.asList(2, "foo")),
+        "d0",
+        "d8",
+        "d9");
+    verifyResults(
+        q.orderBy("array", Query.Direction.DESCENDING).startAfter(Arrays.asList(2, "foo")).limit(2),
+        "d0",
+        "d9");
+    verifyResults(q.orderBy("array").endAt(Collections.singletonList(2)), "d0", "d8", "d9");
+    verifyResults(
+        q.orderBy("array", Query.Direction.DESCENDING).endAt(Collections.singletonList(2)),
+        "d1",
+        "d2");
+    verifyResults(q.orderBy("array").endBefore(Collections.singletonList(2)), "d0", "d8", "d9");
+    ;
+    verifyResults(q.orderBy("array").endBefore(Collections.singletonList(2)).limit(2), "d8", "d9");
+    verifyResults(
+        q.orderBy("array", Query.Direction.DESCENDING).endBefore(Collections.singletonList(2)),
+        "d1",
+        "d2");
+    verifyResults(q.orderBy("array").endBefore(Arrays.asList(2, "foo")), "d0", "d8", "d9");
+    verifyResults(q.orderBy("array").endBefore(Arrays.asList(2, "foo")).limit(2), "d8", "d9");
+    verifyResults(
+        q.orderBy("array", Query.Direction.DESCENDING).endBefore(Arrays.asList(2, "foo")), "d2");
+    verifyResults(q.orderBy("a").orderBy("b").limit(1), "d13");
+    verifyResults(q.orderBy("a", Query.Direction.DESCENDING).orderBy("b").limit(1), "d15");
+    verifyResults(q.orderBy("a").orderBy("b", Query.Direction.DESCENDING).limit(1), "d14");
+    verifyResults(
+        q.orderBy("a", Query.Direction.DESCENDING)
+            .orderBy("b", Query.Direction.DESCENDING)
+            .limit(1),
+        "d16");
+    verifyResults(q.whereEqualTo("null", null), "d18");
+    verifyResults(q.orderBy("null"), "d18");
+    verifyResults(q.whereEqualTo("prefix", Arrays.asList(1, 2)), "d19");
+    verifyResults(q.whereEqualTo("prefix", Arrays.asList(1)).whereEqualTo("suffix", 2), "d20");
+    verifyResults(q.whereEqualTo("map", map()), "d21");
+    verifyResults(q.whereEqualTo("map", map("value", true)), "d22");
+    verifyResults(q.whereEqualTo("map.value", true), "d22");
+    verifyResults(q.orderBy("map"), "d21", "d22", "d23");
+    verifyResults(q.orderBy("map.value"), "d22", "d23");
+  }
+
+  private void verifyResults(Query query, String... documents)
+      throws ExecutionException, InterruptedException {
+    List<QueryDocumentSnapshot> results = query.get().get().getDocuments();
+    List<String> actualKeys =
+        results.stream().map(d -> d.getReference().getId()).collect(Collectors.toList());
+    assertThat(actualKeys).containsExactlyElementsIn(documents);
+  }
+
+  private void addDoc(String key, Map<String, Object> data)
+      throws ExecutionException, InterruptedException {
+    firestore.collection("indexed").document(key).set(data).get();
+  }
+}

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();

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());
+  }
 }

google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITQueryWatchTest.java → google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITWatchTest.java

@@ -35,7 +35,7 @@
 import com.google.cloud.firestore.Query;
 import com.google.cloud.firestore.QueryDocumentSnapshot;
 import com.google.cloud.firestore.QuerySnapshot;
-import com.google.cloud.firestore.it.ITQueryWatchTest.QuerySnapshotEventListener.ListenerAssertions;
+import com.google.cloud.firestore.it.ITWatchTest.QuerySnapshotEventListener.ListenerAssertions;
 import com.google.common.base.Joiner;
 import com.google.common.base.Joiner.MapJoiner;
 import com.google.common.base.Optional;
@@ -59,8 +59,11 @@
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TestName;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
 
-public final class ITQueryWatchTest {
+@RunWith(JUnit4.class)
+public final class ITWatchTest {
 
   private static Firestore firestore;