test: fix flaky session tests (#2913)

mutianf · web-flow · commit a9f8753769c3 · 2026-05-08T08:07:26.000-04:00
Change-Id: I0d0d51c9df31ffd3224866b2345fc1e5c2531676
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java
@@ -182,7 +182,7 @@ void sessionCloseBeforeInit() throws Exception {
   void sessionGoAwayTest() throws Exception {
     SessionImpl session = new SessionImpl(metrics, poolInfo, 0, sessionFactory.createNew());
 
-    Duration goAwayDelay = Duration.ofMillis(100);
+    Duration goAwayDelay = Duration.ofMillis(500);
     FakeSessionListener sessionListener = new FakeSessionListener();
     session.start(
         OpenSessionRequest.newBuilder()
@@ -215,9 +215,14 @@ void sessionGoAwayTest() throws Exception {
       try {
         f.get();
         numOk++;
-      } catch (VRpcException e) {
-        if (e.getResult().getState() == State.UNCOMMITED) {
-          numUncommittedErrors++;
+      } catch (ExecutionException e) {
+        if (e.getCause() instanceof VRpcException) {
+          VRpcException vrpcException = (VRpcException) e.getCause();
+          if (vrpcException.getResult().getState() == State.UNCOMMITED) {
+            numUncommittedErrors++;
+          }
+        } else {
+          throw e;
         }
       }
     }
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java
@@ -77,9 +77,9 @@
 import java.io.IOException;
 import java.time.Duration;
 import java.time.Instant;
-import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
@@ -448,23 +448,44 @@ public void refreshConfigTest() throws Exception {
 
     sessionPool.start(request, new Metadata());
 
-    Thread.sleep(500);
+    long deadline = System.currentTimeMillis() + 10_000;
+    boolean conditionMet = false;
+    List<SessionRequest> requests = null;
+    boolean containsHeader = false;
+
+    while (System.currentTimeMillis() < deadline) {
+      requests = fakeService.getSessionRequests();
+      List<Metadata> headers = headerInterceptor.getHeadersList();
+
+      boolean matchesRefreshRequest = false;
+      for (SessionRequest r : requests) {
+        if (OPEN_SESSION_REQUEST_CORRESPONDENCE.compare(r, refreshRequest)) {
+          matchesRefreshRequest = true;
+          break;
+        }
+      }
+
+      containsHeader = false;
+      for (Metadata header : headers) {
+        if (header.containsKey(metadataKey) && "refresh_value".equals(header.get(metadataKey))) {
+          containsHeader = true;
+          break;
+        }
+      }
 
-    List<SessionRequest> requests = fakeService.getSessionRequests();
+      if (requests.size() > 1 && matchesRefreshRequest && containsHeader) {
+        conditionMet = true;
+        break;
+      }
+
+      Thread.sleep(50);
+    }
+
+    assertThat(conditionMet).isTrue();
     assertThat(requests.size()).isGreaterThan(1);
     assertThat(requests)
         .comparingElementsUsing(OPEN_SESSION_REQUEST_CORRESPONDENCE)
         .contains(refreshRequest);
-
-    // Verify headers
-    List<Metadata> headers = headerInterceptor.getHeadersList();
-    boolean containsHeader = false;
-    for (Metadata header : headers) {
-      if (header.containsKey(metadataKey)) {
-        containsHeader = true;
-        assertThat(header.get(metadataKey)).isEqualTo("refresh_value");
-      }
-    }
     assertThat(containsHeader).isTrue();
   }
 
@@ -519,7 +540,7 @@ public Deadline getObservedDeadline() {
   }
 
   private static class HeaderInterceptor implements ServerInterceptor {
-    private final List<Metadata> headersList = new ArrayList<>();
+    private final List<Metadata> headersList = new CopyOnWriteArrayList<>();
 
     @Override
     public <ReqT, RespT> io.grpc.ServerCall.Listener<ReqT> interceptCall(
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/fake/FakeSessionService.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/fake/FakeSessionService.java
@@ -20,8 +20,8 @@
 import com.google.bigtable.v2.SessionRequest;
 import com.google.bigtable.v2.SessionResponse;
 import io.grpc.stub.StreamObserver;
-import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.atomic.AtomicInteger;
 
@@ -30,7 +30,7 @@ public class FakeSessionService extends FakeSessionGrpc.FakeSessionImplBase {
   private final ScheduledExecutorService executor;
   private final AtomicInteger openRequestCount = new AtomicInteger(0);
 
-  private final List<SessionRequest> sessionRequests = new ArrayList<>();
+  private final List<SessionRequest> sessionRequests = new CopyOnWriteArrayList<>();
 
   public FakeSessionService(ScheduledExecutorService executor) {
     this.executor = executor;
