fix tests

Author: rahul2393

google-cloud-spanner/src/test/java/com/google/cloud/spanner/BatchClientImplTest.java

@@ -57,7 +57,6 @@ public final class BatchClientImplTest {
   private static final String SESSION_NAME = DB_NAME + "/sessions/s1";
   private static final ByteString TXN_ID = ByteString.copyFromUtf8("my-txn");
   private static final String TIMESTAMP = "2017-11-15T10:54:20Z";
-  private static boolean isMultiplexedSession = false;
 
   @Mock private SpannerRpc gapicRpc;
   @Mock private SpannerOptions spannerOptions;
@@ -70,11 +69,6 @@ public final class BatchClientImplTest {
   public static void setupOpenTelemetry() {
     SpannerOptions.resetActiveTracingFramework();
     SpannerOptions.enableOpenTelemetryTraces();
-    Boolean useMultiplexedSessionFromEnvVariablePartitionedOps =
-        SessionPoolOptions.getUseMultiplexedSessionFromEnvVariablePartitionedOps();
-    isMultiplexedSession =
-        useMultiplexedSessionFromEnvVariablePartitionedOps != null
-            && useMultiplexedSessionFromEnvVariablePartitionedOps;
   }
 
   @SuppressWarnings("unchecked")
@@ -95,8 +89,7 @@ public void setUp() {
     when(spannerOptions.getTransportOptions()).thenReturn(transportOptions);
     SessionPoolOptions sessionPoolOptions = mock(SessionPoolOptions.class);
     when(sessionPoolOptions.getPoolMaintainerClock()).thenReturn(Clock.INSTANCE);
-    when(sessionPoolOptions.getUseMultiplexedSessionPartitionedOps())
-        .thenReturn(isMultiplexedSession);
+    when(sessionPoolOptions.getUseMultiplexedSessionPartitionedOps()).thenReturn(true);
     when(sessionPoolOptions.getMultiplexedSessionMaintenanceDuration()).thenReturn(Duration.ZERO);
     when(spannerOptions.getSessionPoolOptions()).thenReturn(sessionPoolOptions);
     @SuppressWarnings("resource")
@@ -108,19 +101,10 @@ public void setUp() {
   @Test
   public void testBatchReadOnlyTxnWithBound() throws Exception {
     Session sessionProto =
-        Session.newBuilder().setName(SESSION_NAME).setMultiplexed(isMultiplexedSession).build();
-    if (isMultiplexedSession) {
-      when(gapicRpc.createSession(
-              eq(DB_NAME),
-              anyString(),
-              anyMap(),
-              optionsCaptor.capture(),
-              eq(isMultiplexedSession)))
-          .thenReturn(sessionProto);
-    } else {
-      when(gapicRpc.createSession(eq(DB_NAME), anyString(), anyMap(), optionsCaptor.capture()))
-          .thenReturn(sessionProto);
-    }
+        Session.newBuilder().setName(SESSION_NAME).setMultiplexed(true).build();
+    when(gapicRpc.createSession(
+            eq(DB_NAME), anyString(), anyMap(), optionsCaptor.capture(), eq(true)))
+        .thenReturn(sessionProto);
     com.google.protobuf.Timestamp timestamp = Timestamps.parse(TIMESTAMP);
     Transaction txnMetadata =
         Transaction.newBuilder().setId(TXN_ID).setReadTimestamp(timestamp).build();

google-cloud-spanner/src/test/java/com/google/cloud/spanner/ChannelUsageTest.java

@@ -23,20 +23,20 @@
 import com.google.cloud.NoCredentials;
 import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime;
 import com.google.cloud.spanner.connection.AbstractMockServerTest;
-import com.google.spanner.v1.BatchCreateSessionsRequest;
+import com.google.spanner.v1.CreateSessionRequest;
 import com.google.spanner.v1.ExecuteSqlRequest;
 import io.grpc.ManagedChannelBuilder;
 import io.grpc.Status;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
-import org.threeten.bp.Duration;
 
 @RunWith(JUnit4.class)
 public class RetryableInternalErrorTest extends AbstractMockServerTest {
   @Test
   public void testTranslateInternalException() {
-    mockSpanner.setBatchCreateSessionsExecutionTime(
+    // With multiplexed sessions, we use CreateSession instead of BatchCreateSessions
+    mockSpanner.setCreateSessionExecutionTime(
         SimulatedExecutionTime.ofException(
             Status.INTERNAL
                 .withDescription("Authentication backend internal server error. Please retry.")
@@ -53,25 +53,18 @@ public void testTranslateInternalException() {
             .setHost(String.format("http://localhost:%d", getPort()))
             .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
             .setCredentials(NoCredentials.getInstance())
-            .setSessionPoolOption(
-                SessionPoolOptions.newBuilder()
-                    .setMinSessions(1)
-                    .setMaxSessions(1)
-                    .setWaitForMinSessions(Duration.ofSeconds(5))
-                    .build())
             .build()
             .getService()) {
 
       DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
-      // Execute a query. This will block until a BatchCreateSessions call has finished and then
+      // Execute a query. This will wait for multiplexed session creation and then
       // invoke ExecuteStreamingSql. Both of these RPCs should be retried.
       try (ResultSet resultSet = client.singleUse().executeQuery(SELECT1_STATEMENT)) {
         assertTrue(resultSet.next());
         assertFalse(resultSet.next());
       }
-      // Verify that both the BatchCreateSessions call and the ExecuteStreamingSql call were
-      // retried.
-      assertEquals(2, mockSpanner.countRequestsOfType(BatchCreateSessionsRequest.class));
+      // Verify that both the CreateSession call and the ExecuteStreamingSql call were retried.
+      assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
       assertEquals(2, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
       // Clear the requests before the next test.
       mockSpanner.clearRequests();