chore(spanner): fix tests in databaseClientimplTest

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

@@ -270,6 +270,9 @@ public void tearDown() {
   @Test
   public void
       testPoolMaintainer_whenInactiveTransactionAndSessionIsNotFoundOnBackend_removeSessionsFromPool() {
+    assumeFalse(
+        "Session pool maintainer test skipped for multiplexed sessions",
+        isMultiplexedSessionsEnabledForRW());
     FakeClock poolMaintainerClock = new FakeClock();
     InactiveTransactionRemovalOptions inactiveTransactionRemovalOptions =
         InactiveTransactionRemovalOptions.newBuilder()
@@ -347,6 +350,9 @@ public void tearDown() {
   @Test
   public void
       testPoolMaintainer_whenInactiveTransactionAndSessionExistsOnBackend_removeSessionsFromPool() {
+    assumeFalse(
+        "Session leaks tests are skipped for multiplexed sessions",
+        isMultiplexedSessionsEnabledForRW());
     FakeClock poolMaintainerClock = new FakeClock();
     InactiveTransactionRemovalOptions inactiveTransactionRemovalOptions =
         InactiveTransactionRemovalOptions.newBuilder()
@@ -482,6 +488,9 @@ public void testPoolMaintainer_whenLongRunningPartitionedUpdateRequest_takeNoAct
    */
   @Test
   public void testPoolMaintainer_whenPDMLFollowedByInactiveTransaction_removeSessionsFromPool() {
+    assumeFalse(
+        "Session leaks tests are skipped for multiplexed sessions",
+        isMultiplexedSessionsEnabledForRW());
     FakeClock poolMaintainerClock = new FakeClock();
     InactiveTransactionRemovalOptions inactiveTransactionRemovalOptions =
         InactiveTransactionRemovalOptions.newBuilder()
@@ -3084,14 +3093,14 @@ public void testDatabaseOrInstanceDoesNotExistOnCreate() {
                         .readWriteTransaction()
                         .run(transaction -> transaction.executeUpdate(UPDATE_STATEMENT)));
             // No additional requests should have been sent by the client.
-            // Note that in case of the use of multiplexed sessions, then we have 2 requests:
+            // Note that in case of the use of regular sessions, then we have 1 request:
             // 1. BatchCreateSessions for the session pool.
-            // 2. CreateSession for the multiplexed session.
-            assertThat(mockSpanner.getRequests())
-                .hasSize(
-                    spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()
-                        ? 2
-                        : 1);
+            // Note that in case of the use of multiplexed sessions for read-only and read-write,
+            // then we have 1 request:
+            // 1. CreateSession for the multiplexed session.
+            // There will be no BatchCreateSessions request in case of multiplexed sessions, because
+            // the session pool options has min size of 0.
+            assertThat(mockSpanner.getRequests()).hasSize(1);
           }
         }
         mockSpanner.reset();
@@ -3211,9 +3220,16 @@ public void testDatabaseOrInstanceIsDeletedAndThenRecreated() throws Exception {
               ResourceNotFoundException.class, () -> dbClient.singleUse().executeQuery(SELECT1));
         }
 
-        assertThrows(
-            ResourceNotFoundException.class,
-            () -> dbClient.readWriteTransaction().run(transaction -> null));
+        if (!spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()) {
+          // We only verify this for read-write transactions if we are not using multiplexed
+          // sessions. For multiplexed sessions, we don't need any special handling, as deleting the
+          // database will also invalidate the multiplexed session, and trying to continue to use it
+          // will continue to return an error.
+          assertThrows(
+              ResourceNotFoundException.class,
+              () -> dbClient.readWriteTransaction().run(transaction -> null));
+        }
+
         assertThat(mockSpanner.getRequests()).isEmpty();
         // Now get a new database client. Normally multiple calls to Spanner#getDatabaseClient will
         // return the same instance, but not when the instance has been invalidated by a
@@ -3300,13 +3316,18 @@ public void testAllowNestedTransactions() throws InterruptedException {
       Thread.sleep(1L);
     }
     assertThat(client.pool.getNumberOfSessionsInPool()).isEqualTo(minSessions);
+    int expectedMinSessions =
+        spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()
+            ? minSessions
+            : minSessions - 1;
     Long res =
         client
             .readWriteTransaction()
             .allowNestedTransaction()
             .run(
                 transaction -> {
-                  assertThat(client.pool.getNumberOfSessionsInPool()).isEqualTo(minSessions - 1);
+                  assertThat(client.pool.getNumberOfSessionsInPool())
+                      .isEqualTo(expectedMinSessions);
                   return transaction.executeUpdate(UPDATE_STATEMENT);
                 });
     assertThat(res).isEqualTo(UPDATE_COUNT);
@@ -3333,6 +3354,9 @@ public void testNestedTransactionsUsingTwoDatabases() throws InterruptedExceptio
     }
     assertThat(client1.pool.getNumberOfSessionsInPool()).isEqualTo(minSessions);
     assertThat(client2.pool.getNumberOfSessionsInPool()).isEqualTo(minSessions);
+    // When read-write transaction uses multiplexed sessions, then sessions are not checked out from
+    // the session pool.
+    int expectedMinSessions = isMultiplexedSessionsEnabledForRW() ? minSessions : minSessions - 1;
     Long res =
         client1
             .readWriteTransaction()
@@ -3341,7 +3365,8 @@ public void testNestedTransactionsUsingTwoDatabases() throws InterruptedExceptio
                 transaction -> {
                   // Client1 should have 1 session checked out.
                   // Client2 should have 0 sessions checked out.
-                  assertThat(client1.pool.getNumberOfSessionsInPool()).isEqualTo(minSessions - 1);
+                  assertThat(client1.pool.getNumberOfSessionsInPool())
+                      .isEqualTo(expectedMinSessions);
                   assertThat(client2.pool.getNumberOfSessionsInPool()).isEqualTo(minSessions);
                   Long add =
                       client2
@@ -3350,9 +3375,9 @@ public void testNestedTransactionsUsingTwoDatabases() throws InterruptedExceptio
                               transaction1 -> {
                                 // Both clients should now have 1 session checked out.
                                 assertThat(client1.pool.getNumberOfSessionsInPool())
-                                    .isEqualTo(minSessions - 1);
+                                    .isEqualTo(expectedMinSessions);
                                 assertThat(client2.pool.getNumberOfSessionsInPool())
-                                    .isEqualTo(minSessions - 1);
+                                    .isEqualTo(expectedMinSessions);
                                 try (ResultSet rs = transaction1.executeQuery(SELECT1)) {
                                   if (rs.next()) {
                                     return rs.getLong(0);
@@ -5090,6 +5115,9 @@ public void testRetryOnResourceExhausted() {
 
   @Test
   public void testSessionPoolExhaustedError_containsStackTraces() {
+    assumeFalse(
+        "Session pool tests are skipped for multiplexed sessions",
+        isMultiplexedSessionsEnabledForRW());
     try (Spanner spanner =
         SpannerOptions.newBuilder()
             .setProjectId(TEST_PROJECT)
@@ -5450,4 +5478,11 @@ private boolean isMultiplexedSessionsEnabled() {
     }
     return spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession();
   }
+
+  private boolean isMultiplexedSessionsEnabledForRW() {
+    if (spanner.getOptions() == null || spanner.getOptions().getSessionPoolOptions() == null) {
+      return false;
+    }
+    return spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW();
+  }
 }