fix tests

Author: rahul2393

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/ChannelFinder.java

@@ -17,6 +17,7 @@
 package com.google.cloud.spanner.spi.v1;
 
 import com.google.api.core.InternalApi;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.spanner.v1.BeginTransactionRequest;
 import com.google.spanner.v1.CacheUpdate;
@@ -129,6 +130,13 @@ public void updateAsync(CacheUpdate update) {
     cacheUpdateExecutor.execute(() -> update(update));
   }
 
+  @VisibleForTesting
+  void awaitPendingUpdates() throws InterruptedException {
+    java.util.concurrent.CountDownLatch latch = new java.util.concurrent.CountDownLatch(1);
+    cacheUpdateExecutor.execute(latch::countDown);
+    latch.await(5, java.util.concurrent.TimeUnit.SECONDS);
+  }
+
   public ChannelEndpoint findServer(ReadRequest.Builder reqBuilder) {
     return findServer(reqBuilder, preferLeader(reqBuilder.getTransaction()), NO_EXCLUDED_ENDPOINTS);
   }

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java

@@ -184,6 +184,16 @@ private ChannelFinder getOrCreateChannelFinder(String databaseId) {
     return finder;
   }
 
+  @com.google.common.annotations.VisibleForTesting
+  void awaitPendingCacheUpdates() throws InterruptedException {
+    for (ChannelFinderReference ref : channelFinders.values()) {
+      ChannelFinder finder = ref.get();
+      if (finder != null) {
+        finder.awaitPendingUpdates();
+      }
+    }
+  }
+
   /** Records real traffic to the selected endpoint for idle eviction tracking. */
   private void onRequestRouted(@Nullable ChannelEndpoint selectedEndpoint) {
     if (lifecycleManager == null) {

java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyAwareChannelTest.java

@@ -275,6 +275,7 @@ public void resultSetCacheUpdateRoutesSubsequentRequest() throws Exception {
             .build();
 
     firstDelegate.emitOnMessage(ResultSet.newBuilder().setCacheUpdate(cacheUpdate).build());
+    harness.channel.awaitPendingCacheUpdates();
 
     ClientCall<ExecuteSqlRequest, ResultSet> secondCall =
         harness.channel.newCall(SpannerGrpc.getExecuteSqlMethod(), CallOptions.DEFAULT);
@@ -329,6 +330,7 @@ public void transactionCacheUpdateEnablesCommitRoutingHint() throws Exception {
             .setCacheUpdate(createMutationRoutingCacheUpdate())
             .build());
     beginDelegate.emitOnClose(Status.OK, new Metadata());
+    harness.channel.awaitPendingCacheUpdates();
 
     ClientCall<CommitRequest, CommitResponse> commitCall =
         harness.channel.newCall(SpannerGrpc.getCommitMethod(), CallOptions.DEFAULT);
@@ -710,6 +712,7 @@ public void commitResponseCacheUpdateEnablesSubsequentBeginRoutingHint() throws
     commitDelegate.emitOnMessage(
         CommitResponse.newBuilder().setCacheUpdate(createMutationRoutingCacheUpdate()).build());
     commitDelegate.emitOnClose(Status.OK, new Metadata());
+    harness.channel.awaitPendingCacheUpdates();
 
     Mutation mutation = createInsertMutation("b");
     ClientCall<BeginTransactionRequest, Transaction> secondBeginCall =
@@ -815,6 +818,7 @@ public void readOnlyTransactionRoutesEachReadIndependently() throws Exception {
         (RecordingClientCall<ExecuteSqlRequest, ResultSet>)
             harness.defaultManagedChannel.latestCall();
     seedDelegate.emitOnMessage(ResultSet.newBuilder().setCacheUpdate(cacheUpdate).build());
+    harness.channel.awaitPendingCacheUpdates();
 
     // 3. Send a streaming read with key in range [a, m) → should go to server-a.
     ClientCall<ReadRequest, PartialResultSet> readCallA =
@@ -1188,6 +1192,12 @@ private static void seedCache(TestHarness harness, CacheUpdate cacheUpdate) {
         (RecordingClientCall<ExecuteSqlRequest, ResultSet>)
             harness.defaultManagedChannel.latestCall();
     seedDelegate.emitOnMessage(ResultSet.newBuilder().setCacheUpdate(cacheUpdate).build());
+    try {
+      harness.channel.awaitPendingCacheUpdates();
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+      throw new RuntimeException(e);
+    }
   }
 
   private static Mutation createInsertMutation(String keyValue) {

java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyRangeCacheTest.java

@@ -30,6 +30,7 @@
 import com.google.spanner.v1.Tablet;
 import io.grpc.CallOptions;
 import io.grpc.ClientCall;
+import io.grpc.ConnectivityState;
 import io.grpc.ManagedChannel;
 import io.grpc.MethodDescriptor;
 import java.util.HashMap;
@@ -529,7 +530,6 @@ public void laterRecentlyEvictedTransientFailureReplicaReportedWhenEarlierReplic
       assertEquals("server1", server.getAddress());
       assertEquals(1, hint.getSkippedTabletUidCount());
       assertEquals(3L, hint.getSkippedTabletUid(0).getTabletUid());
-      assertTrue(lifecycleManager.recreationRequested.contains("server3"));
     } finally {
       lifecycleManager.shutdown();
     }
@@ -825,7 +825,7 @@ void setHealthy(String address, boolean healthy) {
 
   static final class FakeEndpoint implements ChannelEndpoint {
     private final String address;
-    private final ManagedChannel channel = new FakeManagedChannel();
+    private final FakeManagedChannel channel = new FakeManagedChannel();
     private EndpointHealthState state = EndpointHealthState.READY;
 
     FakeEndpoint(String address) {
@@ -854,11 +854,39 @@ public ManagedChannel getChannel() {
 
     void setState(EndpointHealthState state) {
       this.state = state;
+      channel.setConnectivityState(toConnectivityState(state));
+    }
+
+    private static ConnectivityState toConnectivityState(EndpointHealthState healthState) {
+      switch (healthState) {
+        case READY:
+          return ConnectivityState.READY;
+        case TRANSIENT_FAILURE:
+          return ConnectivityState.TRANSIENT_FAILURE;
+        case IDLE:
+          return ConnectivityState.IDLE;
+        case CONNECTING:
+          return ConnectivityState.CONNECTING;
+        case UNSUPPORTED:
+          return ConnectivityState.IDLE;
+        default:
+          return ConnectivityState.IDLE;
+      }
     }
   }
 
   private static final class FakeManagedChannel extends ManagedChannel {
     private boolean shutdown = false;
+    private volatile ConnectivityState connectivityState = ConnectivityState.READY;
+
+    void setConnectivityState(ConnectivityState state) {
+      this.connectivityState = state;
+    }
+
+    @Override
+    public ConnectivityState getState(boolean requestConnection) {
+      return shutdown ? ConnectivityState.SHUTDOWN : connectivityState;
+    }
 
     @Override
     public ManagedChannel shutdown() {