Skip to content

Commit 6cd2a45

Browse files
fix: set timeouts for BatchGetDocuments/RunQuery (#799)
1 parent ec07294 commit 6cd2a45

File tree

5 files changed

+61
-1
lines changed

5 files changed

+61
-1
lines changed

google-cloud-firestore/src/main/java/com/google/cloud/firestore/spi/v1/GrpcFirestoreRpc.java

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,9 @@ public GrpcFirestoreRpc(final FirestoreOptions options) throws IOException {
132132
FirestoreStubSettings.Builder firestoreBuilder =
133133
FirestoreStubSettings.newBuilder(clientContext)
134134
.applyToAllUnaryMethods(retrySettingsSetter);
135+
// Manually apply the retry settings to streaming methods
136+
firestoreBuilder.runQuerySettings().setRetrySettings(options.getRetrySettings());
137+
firestoreBuilder.batchGetDocumentsSettings().setRetrySettings(options.getRetrySettings());
135138
firestoreStub = GrpcFirestoreStub.create(firestoreBuilder.build());
136139
} catch (Exception e) {
137140
throw new IOException(e);

google-cloud-firestore/src/test/java/com/google/cloud/firestore/ITBulkWriterTest.java renamed to google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITBulkWriterTest.java

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
* limitations under the License.
1515
*/
1616

17-
package com.google.cloud.firestore;
17+
package com.google.cloud.firestore.it;
1818

1919
import static org.junit.Assert.assertEquals;
2020
import static org.junit.Assert.assertNotNull;
@@ -24,6 +24,14 @@
2424

2525
import com.google.api.core.ApiFuture;
2626
import com.google.cloud.Timestamp;
27+
import com.google.cloud.firestore.BulkWriter;
28+
import com.google.cloud.firestore.CollectionReference;
29+
import com.google.cloud.firestore.DocumentReference;
30+
import com.google.cloud.firestore.DocumentSnapshot;
31+
import com.google.cloud.firestore.Firestore;
32+
import com.google.cloud.firestore.FirestoreOptions;
33+
import com.google.cloud.firestore.LocalFirestoreHelper;
34+
import com.google.cloud.firestore.WriteResult;
2735
import com.google.common.base.Preconditions;
2836
import java.util.ArrayList;
2937
import java.util.Collections;

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,10 @@
5858
import org.junit.Rule;
5959
import org.junit.Test;
6060
import org.junit.rules.TestName;
61+
import org.junit.runner.RunWith;
62+
import org.junit.runners.JUnit4;
6163

64+
@RunWith(JUnit4.class)
6265
public final class ITQueryWatchTest {
6366

6467
private static Firestore firestore;

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,10 @@
3232
import org.junit.Test;
3333
import org.junit.rules.TestName;
3434
import org.junit.rules.Timeout;
35+
import org.junit.runner.RunWith;
36+
import org.junit.runners.JUnit4;
3537

38+
@RunWith(JUnit4.class)
3639
public class ITShutdownTest {
3740
@Rule public final Timeout timeout = new Timeout(5, TimeUnit.SECONDS);
3841
@Rule public TestName testName = new TestName();

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

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,12 +33,14 @@
3333
import static org.junit.Assert.assertNotEquals;
3434
import static org.junit.Assert.assertNotNull;
3535
import static org.junit.Assert.assertNull;
36+
import static org.junit.Assert.assertThrows;
3637
import static org.junit.Assert.assertTrue;
3738
import static org.junit.Assert.fail;
3839

3940
import com.google.api.core.ApiFuture;
4041
import com.google.api.core.ApiFutures;
4142
import com.google.api.core.SettableApiFuture;
43+
import com.google.api.gax.retrying.RetrySettings;
4244
import com.google.api.gax.rpc.ApiStreamObserver;
4345
import com.google.cloud.Timestamp;
4446
import com.google.cloud.firestore.BulkWriter;
@@ -100,7 +102,11 @@
100102
import org.junit.Rule;
101103
import org.junit.Test;
102104
import org.junit.rules.TestName;
105+
import org.junit.runner.RunWith;
106+
import org.junit.runners.JUnit4;
107+
import org.threeten.bp.Duration;
103108

109+
@RunWith(JUnit4.class)
104110
public class ITSystemTest {
105111

106112
private static final double DOUBLE_EPSILON = 0.000001;
@@ -1761,4 +1767,41 @@ public void testRecursiveDeleteWithCustomBulkWriterInstance() throws Exception {
17611767
assertEquals(0, countCollectionChildren(randomColl));
17621768
assertEquals(6, callbackCount[0]);
17631769
}
1770+
1771+
@Test
1772+
public void testEnforcesTimeouts() throws Exception {
1773+
FirestoreOptions firestoreOptions =
1774+
FirestoreOptions.newBuilder()
1775+
.setRetrySettings(
1776+
RetrySettings.newBuilder()
1777+
.setMaxRpcTimeout(Duration.ofMillis(1))
1778+
.setTotalTimeout(Duration.ofMillis(1))
1779+
.setInitialRpcTimeout(Duration.ofMillis(1))
1780+
.build())
1781+
.build();
1782+
firestore = firestoreOptions.getService();
1783+
CollectionReference collection = firestore.collection("timeout");
1784+
1785+
// RunQuery
1786+
assertThrows(ExecutionException.class, () -> collection.get().get());
1787+
// CommitRequest
1788+
assertThrows(ExecutionException.class, () -> collection.add(map()).get());
1789+
// BulkCommit
1790+
assertThrows(
1791+
ExecutionException.class,
1792+
() -> {
1793+
BulkWriter bulkWriter = firestore.bulkWriter();
1794+
ApiFuture<WriteResult> op = bulkWriter.set(collection.document(), map());
1795+
bulkWriter.close();
1796+
op.get();
1797+
});
1798+
// BatchGetDocuments
1799+
assertThrows(ExecutionException.class, () -> collection.document().get().get());
1800+
// ListDocuments
1801+
assertThrows(FirestoreException.class, () -> collection.listDocuments().iterator().hasNext());
1802+
// ListCollections
1803+
assertThrows(
1804+
FirestoreException.class,
1805+
() -> collection.document().listCollections().iterator().hasNext());
1806+
}
17641807
}

0 commit comments

Comments
 (0)