test: add RetryInfo test to test server disabled sending retry info (#…

…2048) Add more retry info test to test the scenario where server stopped sending retry info and client handles retry correctly. Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly: - [ ] Make sure to open an issue as a [bug/issue](https://togithub.com/googleapis/java-bigtable/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [ ] Ensure the tests and linter pass - [ ] Code coverage does not decrease (if any source code was changed) - [ ] Appropriate docs were updated (if necessary) - [ ] Rollback plan is reviewed and LGTMed Fixes #<issue_number_goes_here> ☕️ If you write sample code, please follow the [samples format]( https://togithub.com/GoogleCloudPlatform/java-docs-samples/blob/main/SAMPLE_FORMAT.md).
googleapis · Jan 11, 2024 · 55d1fb9 · 55d1fb9
1 parent 94d1ea0
commit 55d1fb9
Showing 1 changed file with 247 additions and 9 deletions.
diff --git a/...le-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/RetryInfoTest.java b/...le-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/RetryInfoTest.java
@@ -16,6 +16,7 @@
 package com.google.cloud.bigtable.data.v2.stub;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertThrows;
 
 import com.google.api.gax.core.NoCredentialsProvider;
 import com.google.api.gax.grpc.GrpcStatusCode;
@@ -132,6 +133,20 @@ public void testReadRowDisableRetryInfo() throws IOException {
     }
   }
 
+  @Test
+  public void testReadRowServerNotReturningRetryInfo() {
+    verifyNoRetryInfo(() -> client.readRow("table", "row"), true);
+  }
+
+  @Test
+  public void testReadRowServerNotReturningRetryInfoClientDisabledHandling() throws IOException {
+    settings.stubSettings().setEnableRetryInfo(false);
+
+    try (BigtableDataClient newClient = BigtableDataClient.create(settings.build())) {
+      verifyNoRetryInfo(() -> newClient.readRow("table", "row"), true);
+    }
+  }
+
   @Test
   public void testReadRows() {
     verifyRetryInfoIsUsed(() -> client.readRows(Query.create("table")).iterator().hasNext(), true);
@@ -152,6 +167,20 @@ public void testReadRowsDisableRetryInfo() throws IOException {
     }
   }
 
+  @Test
+  public void testReadRowsServerNotReturningRetryInfo() {
+    verifyNoRetryInfo(() -> client.readRows(Query.create("table")).iterator().hasNext(), true);
+  }
+
+  @Test
+  public void testReadRowsServerNotReturningRetryInfoClientDisabledHandling() throws IOException {
+    settings.stubSettings().setEnableRetryInfo(false);
+
+    try (BigtableDataClient newClient = BigtableDataClient.create(settings.build())) {
+      verifyNoRetryInfo(() -> newClient.readRows(Query.create("table")).iterator().hasNext(), true);
+    }
+  }
+
   @Test
   public void testMutateRows() {
     verifyRetryInfoIsUsed(
@@ -185,6 +214,30 @@ public void testMutateRowsDisableRetryInfo() throws IOException {
     }
   }
 
+  @Test
+  public void testMutateRowsServerNotReturningRetryInfo() {
+    verifyNoRetryInfo(
+        () ->
+            client.bulkMutateRows(
+                BulkMutation.create("fake-table")
+                    .add(RowMutationEntry.create("row-key-1").setCell("cf", "q", "v"))),
+        true);
+  }
+
+  @Test
+  public void testMutateRowsServerNotReturningRetryInfoClientDisabledHandling() throws IOException {
+    settings.stubSettings().setEnableRetryInfo(false);
+
+    try (BigtableDataClient newClient = BigtableDataClient.create(settings.build())) {
+      verifyNoRetryInfo(
+          () ->
+              newClient.bulkMutateRows(
+                  BulkMutation.create("fake-table")
+                      .add(RowMutationEntry.create("row-key-1").setCell("cf", "q", "v"))),
+          true);
+    }
+  }
+
   @Test
   public void testMutateRow() {
     verifyRetryInfoIsUsed(
@@ -208,6 +261,23 @@ public void testMutateRowDisableRetryInfo() throws IOException {
     }
   }
 
+  @Test
+  public void testMutateRowServerNotReturningRetryInfo() {
+    verifyNoRetryInfo(
+        () -> client.mutateRow(RowMutation.create("table", "key").setCell("cf", "q", "v")), true);
+  }
+
+  @Test
+  public void testMutateRowServerNotReturningRetryInfoClientDisabledHandling() throws IOException {
+    settings.stubSettings().setEnableRetryInfo(false);
+
+    try (BigtableDataClient newClient = BigtableDataClient.create(settings.build())) {
+      verifyNoRetryInfo(
+          () -> newClient.mutateRow(RowMutation.create("table", "key").setCell("cf", "q", "v")),
+          true);
+    }
+  }
+
   @Test
   public void testSampleRowKeys() {
     verifyRetryInfoIsUsed(() -> client.sampleRowKeys("table"), true);
@@ -227,6 +297,21 @@ public void testSampleRowKeysDisableRetryInfo() throws IOException {
     }
   }
 
+  @Test
+  public void testSampleRowKeysServerNotReturningRetryInfo() {
+    verifyNoRetryInfo(() -> client.sampleRowKeys("table"), true);
+  }
+
+  @Test
+  public void testSampleRowKeysServerNotReturningRetryInfoClientDisabledHandling()
+      throws IOException {
+    settings.stubSettings().setEnableRetryInfo(false);
+
+    try (BigtableDataClient newClient = BigtableDataClient.create(settings.build())) {
+      verifyNoRetryInfo(() -> newClient.sampleRowKeys("table"), true);
+    }
+  }
+
   @Test
   public void testCheckAndMutateRow() {
     verifyRetryInfoIsUsed(
@@ -256,6 +341,33 @@ public void testCheckAndMutateDisableRetryInfo() throws IOException {
     }
   }
 
+  @Test
+  public void testCheckAndMutateServerNotReturningRetryInfo() {
+    verifyNoRetryInfo(
+        () ->
+            client.checkAndMutateRow(
+                ConditionalRowMutation.create("table", "key")
+                    .condition(Filters.FILTERS.value().regex("old-value"))
+                    .then(Mutation.create().setCell("cf", "q", "v"))),
+        false);
+  }
+
+  @Test
+  public void testCheckAndMutateServerNotReturningRetryInfoClientDisabledHandling()
+      throws IOException {
+    settings.stubSettings().setEnableRetryInfo(false);
+
+    try (BigtableDataClient newClient = BigtableDataClient.create(settings.build())) {
+      verifyNoRetryInfo(
+          () ->
+              newClient.checkAndMutateRow(
+                  ConditionalRowMutation.create("table", "key")
+                      .condition(Filters.FILTERS.value().regex("old-value"))
+                      .then(Mutation.create().setCell("cf", "q", "v"))),
+          false);
+    }
+  }
+
   @Test
   public void testReadModifyWrite() {
     verifyRetryInfoIsUsed(
@@ -280,6 +392,28 @@ public void testReadModifyWriteDisableRetryInfo() throws IOException {
     }
   }
 
+  @Test
+  public void testReadModifyWriteServerNotReturningRetryInfo() {
+    verifyNoRetryInfo(
+        () ->
+            client.readModifyWriteRow(
+                ReadModifyWriteRow.create("table", "row").append("cf", "q", "v")),
+        false);
+  }
+
+  @Test
+  public void testReadModifyWriteNotReturningRetryInfoClientDisabledHandling() throws IOException {
+    settings.stubSettings().setEnableRetryInfo(false);
+
+    try (BigtableDataClient newClient = BigtableDataClient.create(settings.build())) {
+      verifyNoRetryInfo(
+          () ->
+              newClient.readModifyWriteRow(
+                  ReadModifyWriteRow.create("table", "row").append("cf", "q", "v")),
+          false);
+    }
+  }
+
   @Test
   public void testReadChangeStream() {
     verifyRetryInfoIsUsed(
@@ -308,6 +442,28 @@ public void testReadChangeStreamDisableRetryInfo() throws IOException {
     }
   }
 
+  @Test
+  public void testReadChangeStreamServerNotReturningRetryInfo() {
+    verifyNoRetryInfo(
+        () -> client.readChangeStream(ReadChangeStreamQuery.create("table")).iterator().hasNext(),
+        true);
+  }
+
+  @Test
+  public void testReadChangeStreamNotReturningRetryInfoClientDisabledHandling() throws IOException {
+    settings.stubSettings().setEnableRetryInfo(false);
+
+    try (BigtableDataClient newClient = BigtableDataClient.create(settings.build())) {
+      verifyNoRetryInfo(
+          () ->
+              newClient
+                  .readChangeStream(ReadChangeStreamQuery.create("table"))
+                  .iterator()
+                  .hasNext(),
+          true);
+    }
+  }
+
   @Test
   public void testGenerateInitialChangeStreamPartition() {
     verifyRetryInfoIsUsed(
@@ -330,6 +486,25 @@ public void testGenerateInitialChangeStreamPartitionDisableRetryInfo() throws IO
     }
   }
 
+  @Test
+  public void testGenerateInitialChangeStreamServerNotReturningRetryInfo() {
+    verifyNoRetryInfo(
+        () -> client.generateInitialChangeStreamPartitions("table").iterator().hasNext(), true);
+  }
+
+  @Test
+  public void testGenerateInitialChangeStreamServerNotReturningRetryInfoClientDisabledHandling()
+      throws IOException {
+    settings.stubSettings().setEnableRetryInfo(false);
+
+    try (BigtableDataClient newClient = BigtableDataClient.create(settings.build())) {
+      verifyNoRetryInfo(
+          () -> newClient.generateInitialChangeStreamPartitions("table").iterator().hasNext(),
+          true);
+    }
+  }
+
+  // Test the case where server returns retry info and client enables handling of retry info
   private void verifyRetryInfoIsUsed(Runnable runnable, boolean retryableError) {
     if (retryableError) {
       enqueueRetryableExceptionWithDelay(delay);
@@ -344,6 +519,7 @@ private void verifyRetryInfoIsUsed(Runnable runnable, boolean retryableError) {
     assertThat(stopwatch.elapsed()).isAtLeast(Duration.ofSeconds(delay.getSeconds()));
   }
 
+  // Test the case where server returns retry info but client disabled handling of retry info
   private void verifyRetryInfoCanBeDisabled(Runnable runnable) {
     enqueueRetryableExceptionWithDelay(delay);
     Stopwatch stopwatch = Stopwatch.createStarted();
@@ -354,17 +530,58 @@ private void verifyRetryInfoCanBeDisabled(Runnable runnable) {
     assertThat(stopwatch.elapsed()).isLessThan(Duration.ofSeconds(delay.getSeconds()));
 
     attemptCounter.set(0);
-    ApiException exception = enqueueNonRetryableExceptionWithDelay(delay);
-    try {
+    ApiException expectedApiException = enqueueNonRetryableExceptionWithDelay(delay);
+    ApiException actualException =
+        assertThrows("non retryable operations should fail", ApiException.class, runnable::run);
+    if (actualException instanceof MutateRowsException) {
+      assertThat(
+              ((MutateRowsException) actualException)
+                  .getFailedMutations()
+                  .get(0)
+                  .getError()
+                  .getStatusCode())
+          .isEqualTo(expectedApiException.getStatusCode());
+    } else {
+      assertThat(actualException.getStatusCode()).isEqualTo(expectedApiException.getStatusCode());
+    }
+    assertThat(attemptCounter.get()).isEqualTo(1);
+  }
+
+  // Test the case where server does not return retry info
+  private void verifyNoRetryInfo(Runnable runnable, boolean operationRetryable) {
+    enqueueRetryableExceptionNoRetryInfo();
+
+    if (!operationRetryable) {
+      assertThrows("non retryable operation should fail", ApiException.class, runnable::run);
+      assertThat(attemptCounter.get()).isEqualTo(1);
+    } else {
+      Stopwatch stopwatch = Stopwatch.createStarted();
       runnable.run();
-    } catch (ApiException e) {
-      if (e instanceof MutateRowsException) {
-        assertThat(((MutateRowsException) e).getFailedMutations().get(0).getError().getStatusCode())
-            .isEqualTo(exception.getStatusCode());
-      } else {
-        assertThat(e.getStatusCode()).isEqualTo(exception.getStatusCode());
-      }
+      stopwatch.stop();
+
+      assertThat(attemptCounter.get()).isEqualTo(2);
+      assertThat(stopwatch.elapsed()).isLessThan(Duration.ofSeconds(delay.getSeconds()));
     }
+
+    attemptCounter.set(0);
+
+    ApiException expectedApiException = enqueueNonRetryableExceptionNoRetryInfo();
+
+    ApiException actualApiException =
+        assertThrows("non retryable error should fail", ApiException.class, runnable::run);
+    if (actualApiException instanceof MutateRowsException) {
+      assertThat(
+              ((MutateRowsException) actualApiException)
+                  .getFailedMutations()
+                  .get(0)
+                  .getError()
+                  .getStatusCode())
+          .isEqualTo(expectedApiException.getStatusCode());
+    } else {
+      assertThat(actualApiException.getStatusCode())
+          .isEqualTo(expectedApiException.getStatusCode());
+    }
+
     assertThat(attemptCounter.get()).isEqualTo(1);
   }
 
@@ -408,6 +625,27 @@ private ApiException enqueueNonRetryableExceptionWithDelay(com.google.protobuf.D
     return exception;
   }
 
+  private void enqueueRetryableExceptionNoRetryInfo() {
+    ApiException exception =
+        new UnavailableException(
+            new StatusRuntimeException(Status.UNAVAILABLE),
+            GrpcStatusCode.of(Status.Code.UNAVAILABLE),
+            true);
+    service.expectations.add(exception);
+  }
+
+  private ApiException enqueueNonRetryableExceptionNoRetryInfo() {
+    ApiException exception =
+        new InternalException(
+            new StatusRuntimeException(Status.INTERNAL),
+            GrpcStatusCode.of(Status.Code.INTERNAL),
+            false);
+
+    service.expectations.add(exception);
+
+    return exception;
+  }
+
   private class FakeBigtableService extends BigtableGrpc.BigtableImplBase {
     Queue<Exception> expectations = Queues.newArrayDeque();