From fed3041b976fb6844921ae266fee898c687208f7 Mon Sep 17 00:00:00 2001
From: Sri Harsha CH <sriharshach@google.com>
Date: Thu, 19 Dec 2024 05:49:06 +0000
Subject: [PATCH] chore(spanner): fix tests in SessionPoolLeakTest

---
 .../MultiplexedSessionDatabaseClient.java     |  2 +-
 .../cloud/spanner/SessionPoolLeakTest.java    | 19 +++++++++++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java
index b8abf978536..aed7c3f3afb 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java
@@ -262,7 +262,7 @@ public void onSessionReady(SessionImpl session) {
                     .getOptions()
                     .getSessionPoolOptions()
                     .getSkipVerifyBeginTransactionForMuxRW()) {
-              
+
               verifyBeginTransactionWithRWOnMultiplexedSessionAsync(session.getName());
             }
           }
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolLeakTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolLeakTest.java
index 4672f03aeff..8ccea443dc1 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolLeakTest.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolLeakTest.java
@@ -171,6 +171,9 @@ public void testIgnoreLeakedSession() {
 
   @Test
   public void testReadWriteTransactionExceptionOnCreateSession() {
+    assumeFalse(
+        "Session Leaks do not occur with Multiplexed Sessions",
+        isMultiplexedSessionsEnabledForRW());
     readWriteTransactionTest(
         () ->
             mockSpanner.setBatchCreateSessionsExecutionTime(
@@ -180,6 +183,9 @@ public void testReadWriteTransactionExceptionOnCreateSession() {
 
   @Test
   public void testReadWriteTransactionExceptionOnBegin() {
+    assumeFalse(
+        "Session Leaks do not occur with Multiplexed Sessions",
+        isMultiplexedSessionsEnabledForRW());
     readWriteTransactionTest(
         () ->
             mockSpanner.setBeginTransactionExecutionTime(
@@ -200,6 +206,9 @@ private void readWriteTransactionTest(
 
   @Test
   public void testTransactionManagerExceptionOnCreateSession() {
+    assumeFalse(
+        "Session Leaks do not occur with Multiplexed Sessions",
+        isMultiplexedSessionsEnabledForRW());
     transactionManagerTest(
         () ->
             mockSpanner.setBatchCreateSessionsExecutionTime(
@@ -209,6 +218,9 @@ public void testTransactionManagerExceptionOnCreateSession() {
 
   @Test
   public void testTransactionManagerExceptionOnBegin() {
+    assumeFalse(
+        "Session Leaks do not occur with Multiplexed Sessions",
+        isMultiplexedSessionsEnabledForRW());
     assertThat(pool.getNumberOfSessionsInPool(), is(equalTo(0)));
     mockSpanner.setBeginTransactionExecutionTime(
         SimulatedExecutionTime.ofException(FAILED_PRECONDITION));
@@ -229,4 +241,11 @@ private void transactionManagerTest(Runnable setup, int expectedNumberOfSessions
     }
     assertEquals(expectedNumberOfSessionsAfterExecution, pool.getNumberOfSessionsInPool());
   }
+
+  private boolean isMultiplexedSessionsEnabledForRW() {
+    if (spanner.getOptions() == null || spanner.getOptions().getSessionPoolOptions() == null) {
+      return false;
+    }
+    return spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW();
+  }
 }