test: fix flaky vip_java_bigtable test (#2116)

googleapis · Feb 23, 2024 · 899c50b · 899c50b
1 parent 4d70dee
commit 899c50b
Showing 1 changed file with 17 additions and 10 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
@@ -88,7 +88,7 @@ public class RetryInfoTest {
   private BigtableDataSettings.Builder settings;
 
   private AtomicInteger attemptCounter = new AtomicInteger();
-  private com.google.protobuf.Duration delay =
+  private com.google.protobuf.Duration defaultDelay =
       com.google.protobuf.Duration.newBuilder().setSeconds(2).setNanos(0).build();
 
   @Before
@@ -328,7 +328,7 @@ public void testCheckAndMutateDisableRetryInfo() throws IOException {
     settings.stubSettings().setEnableRetryInfo(false);
 
     try (BigtableDataClient client = BigtableDataClient.create(settings.build())) {
-      ApiException exception = enqueueNonRetryableExceptionWithDelay(delay);
+      ApiException exception = enqueueNonRetryableExceptionWithDelay(defaultDelay);
       try {
         client.checkAndMutateRow(
             ConditionalRowMutation.create("table", "key")
@@ -382,7 +382,7 @@ public void testReadModifyWriteDisableRetryInfo() throws IOException {
     settings.stubSettings().setEnableRetryInfo(false);
 
     try (BigtableDataClient client = BigtableDataClient.create(settings.build())) {
-      ApiException exception = enqueueNonRetryableExceptionWithDelay(delay);
+      ApiException exception = enqueueNonRetryableExceptionWithDelay(defaultDelay);
       try {
         client.readModifyWriteRow(ReadModifyWriteRow.create("table", "row").append("cf", "q", "v"));
       } catch (ApiException e) {
@@ -460,7 +460,8 @@ public void testReadChangeStreamNotReturningRetryInfoClientDisabledHandling() th
                   .readChangeStream(ReadChangeStreamQuery.create("table"))
                   .iterator()
                   .hasNext(),
-          true);
+          true,
+          com.google.protobuf.Duration.newBuilder().setSeconds(5).setNanos(0).build());
     }
   }
 
@@ -507,30 +508,30 @@ public void testGenerateInitialChangeStreamServerNotReturningRetryInfoClientDisa
   // 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);
+      enqueueRetryableExceptionWithDelay(defaultDelay);
     } else {
-      enqueueNonRetryableExceptionWithDelay(delay);
+      enqueueNonRetryableExceptionWithDelay(defaultDelay);
     }
     Stopwatch stopwatch = Stopwatch.createStarted();
     runnable.run();
     stopwatch.stop();
 
     assertThat(attemptCounter.get()).isEqualTo(2);
-    assertThat(stopwatch.elapsed()).isAtLeast(Duration.ofSeconds(delay.getSeconds()));
+    assertThat(stopwatch.elapsed()).isAtLeast(Duration.ofSeconds(defaultDelay.getSeconds()));
   }
 
   // Test the case where server returns retry info but client disabled handling of retry info
   private void verifyRetryInfoCanBeDisabled(Runnable runnable) {
-    enqueueRetryableExceptionWithDelay(delay);
+    enqueueRetryableExceptionWithDelay(defaultDelay);
     Stopwatch stopwatch = Stopwatch.createStarted();
     runnable.run();
     stopwatch.stop();
 
     assertThat(attemptCounter.get()).isEqualTo(2);
-    assertThat(stopwatch.elapsed()).isLessThan(Duration.ofSeconds(delay.getSeconds()));
+    assertThat(stopwatch.elapsed()).isLessThan(Duration.ofSeconds(defaultDelay.getSeconds()));
 
     attemptCounter.set(0);
-    ApiException expectedApiException = enqueueNonRetryableExceptionWithDelay(delay);
+    ApiException expectedApiException = enqueueNonRetryableExceptionWithDelay(defaultDelay);
     ApiException actualException =
         assertThrows("non retryable operations should fail", ApiException.class, runnable::run);
     if (actualException instanceof MutateRowsException) {
@@ -549,6 +550,12 @@ private void verifyRetryInfoCanBeDisabled(Runnable runnable) {
 
   // Test the case where server does not return retry info
   private void verifyNoRetryInfo(Runnable runnable, boolean operationRetryable) {
+    verifyNoRetryInfo(runnable, operationRetryable, defaultDelay);
+  }
+
+  // individual test can override the default delay
+  private void verifyNoRetryInfo(
+      Runnable runnable, boolean operationRetryable, com.google.protobuf.Duration delay) {
     enqueueRetryableExceptionNoRetryInfo();
 
     if (!operationRetryable) {