tests: add test for retry on UNAVAILABLE

googleapis · May 14, 2020 · aac7cf7 · aac7cf7
1 parent afbb2c4
commit aac7cf7
Showing 1 changed file with 17 additions and 0 deletions.
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java
@@ -324,6 +324,23 @@ public Long run(TransactionContext transaction) throws Exception {
     }
   }
 
+  @Test
+  public void testPartitionedDmlRetriesOnUnavailable() throws Exception {
+    mockSpanner.setExecuteSqlExecutionTime(
+        SimulatedExecutionTime.ofException(Status.UNAVAILABLE.asRuntimeException()));
+    SpannerOptions.Builder builder =
+        SpannerOptions.newBuilder()
+            .setProjectId(TEST_PROJECT)
+            .setChannelProvider(channelProvider)
+            .setCredentials(NoCredentials.getInstance());
+    try (Spanner spanner = builder.build().getService()) {
+      DatabaseClient client =
+          spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
+      long updateCount = client.executePartitionedUpdate(UPDATE_STATEMENT);
+      assertThat(updateCount).isEqualTo(UPDATE_COUNT);
+    }
+  }
+
   @Test
   public void testDatabaseOrInstanceDoesNotExistOnPrepareSession() throws Exception {
     StatusRuntimeException[] exceptions =