From 00e2a1cfaa33a59a91829266853624a46377e5c0 Mon Sep 17 00:00:00 2001
From: Jack Dingilian <jackdingilian@google.com>
Date: Thu, 26 Sep 2024 12:52:23 -0400
Subject: [PATCH] fix: time based flakiness in execute query deadline test

Change-Id: I93c1c03a0c41c92dbe65b5ec4888e7df526ad457
---
 .../v2/stub/sql/ExecuteQueryCallableTest.java | 26 ++++++++++++-------
 1 file changed, 17 insertions(+), 9 deletions(-)

diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallableTest.java
index 14275d3cd8..deedfbaba1 100644
--- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallableTest.java
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/sql/ExecuteQueryCallableTest.java
@@ -23,6 +23,7 @@
 import static org.junit.Assert.assertThrows;
 
 import com.google.api.gax.retrying.RetrySettings;
+import com.google.api.gax.rpc.DeadlineExceededException;
 import com.google.api.gax.rpc.UnavailableException;
 import com.google.bigtable.v2.BigtableGrpc;
 import com.google.bigtable.v2.ExecuteQueryRequest;
@@ -36,7 +37,6 @@
 import com.google.cloud.bigtable.data.v2.models.sql.Statement;
 import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStub;
 import com.google.cloud.bigtable.gaxx.testing.FakeStreamingApi.ServerStreamingStashCallable;
-import com.google.common.collect.Range;
 import io.grpc.Context;
 import io.grpc.Deadline;
 import io.grpc.Server;
@@ -141,12 +141,13 @@ public void testExecuteQueryRequestsSetDefaultDeadline() {
     Iterator<SqlRow> iterator = stream.rows().iterator();
     // We don't care about this but are reusing the fake service that tests retries
     assertThrows(UnavailableException.class, iterator::next).getCause();
-    // We have 30s default, we assume less than 1s has been burned when the fake service sets it
-    assertThat(fakeService.deadlineMillisRemaining).isIn(Range.closed(29000L, 30100L));
+    // We have 30s default, we give it a wide range to avoid flakiness, this is mostly just checking
+    // that some default is set
+    assertThat(fakeService.deadlineMillisRemaining).isLessThan(30001L);
   }
 
   @Test
-  public void testExecuteQueryRequestsRespectOverridenDeadline() throws IOException {
+  public void testExecuteQueryRequestsRespectDeadline() throws IOException {
     BigtableDataSettings.Builder overrideSettings =
         BigtableDataSettings.newBuilderForEmulator(server.getPort())
             .setProjectId("fake-project")
@@ -156,8 +157,8 @@ public void testExecuteQueryRequestsRespectOverridenDeadline() throws IOExceptio
         .executeQuerySettings()
         .setRetrySettings(
             RetrySettings.newBuilder()
-                .setInitialRpcTimeout(Duration.ofMinutes(5))
-                .setMaxRpcTimeout(Duration.ofMinutes(5))
+                .setInitialRpcTimeout(Duration.ofMillis(10))
+                .setMaxRpcTimeout(Duration.ofMillis(10))
                 .build());
     EnhancedBigtableStub overrideDeadline =
         EnhancedBigtableStub.create(overrideSettings.build().getStubSettings());
@@ -165,9 +166,7 @@ public void testExecuteQueryRequestsRespectOverridenDeadline() throws IOExceptio
         overrideDeadline.executeQueryCallable().call(Statement.of("SELECT * FROM table"));
     Iterator<SqlRow> overrideIterator = streamOverride.rows().iterator();
     // We don't care about this but are reusing the fake service that tests retries
-    assertThrows(UnavailableException.class, overrideIterator::next).getCause();
-    // We have 30s default, we assume less than 1s has been burned when the fake service sets it
-    assertThat(fakeService.deadlineMillisRemaining).isIn(Range.closed(299000L, 300100L));
+    assertThrows(DeadlineExceededException.class, overrideIterator::next).getCause();
   }
 
   private static class FakeService extends BigtableGrpc.BigtableImplBase {
@@ -181,6 +180,15 @@ public void executeQuery(
       Deadline deadline = Context.current().getDeadline();
       if (deadline != null) {
         deadlineMillisRemaining = deadline.timeRemaining(TimeUnit.MILLISECONDS);
+      } else {
+        // set to max long when deadline isn't set
+        deadlineMillisRemaining = Long.MAX_VALUE;
+      }
+      // Sleep for 100ms to trigger deadline exceeded for tests with a shorter deadline
+      try {
+        Thread.sleep(100);
+      } catch (InterruptedException e) {
+        throw new RuntimeException(e);
       }
       attempts++;
       responseObserver.onNext(metadata(columnMetadata("test", stringType())));