chore(spanner): fix test

Author: harshachinta

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/PartitionedQueryMockServerTest.java

@@ -93,7 +93,13 @@ public void testPartitionQuery() {
           assertFalse(resultSet.next());
         }
       }
-      assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
+      if (isMultiplexedSessionsEnabledForPartitionedOps(connection.getSpanner())) {
+        assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
+      } else if (isMultiplexedSessionsEnabled(connection.getSpanner())) {
+        assertEquals(3, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
+      } else {
+        assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
+      }
     }
     assertEquals(2, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
     assertEquals(2, mockSpanner.countRequestsOfType(PartitionQueryRequest.class));
@@ -151,7 +157,13 @@ public void testMixNormalAndPartitionQueryInReadOnlyTransaction() {
         readTimestamps.add(connection.getReadTimestamp());
         connection.commit();
       }
-      assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
+      if (isMultiplexedSessionsEnabledForPartitionedOps(connection.getSpanner())) {
+        assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
+      } else if (isMultiplexedSessionsEnabled(connection.getSpanner())) {
+        assertEquals(3, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
+      } else {
+        assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
+      }
     }
     // The above will start two transactions:
     // 1. The initial 'normal' read-only transaction.
@@ -570,7 +582,13 @@ public void testRunPartitionedQueryUsingSql() {
           assertEquals(maxPartitions * generatedRowCount, rowCount);
         }
       }
-      assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
+      if (isMultiplexedSessionsEnabledForPartitionedOps(connection.getSpanner())) {
+        assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
+      } else if (isMultiplexedSessionsEnabled(connection.getSpanner())) {
+        assertEquals(3, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
+      } else {
+        assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
+      }
     }
     // We have 2 requests of each, as we run the query with data boost both enabled and disabled.
     assertEquals(2, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
@@ -755,7 +773,13 @@ public void testAutoPartitionMode() {
               .getMessage()
               .contains("Partition query is not supported for read/write transaction"));
 
-      assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
+      if (isMultiplexedSessionsEnabledForPartitionedOps(connection.getSpanner())) {
+        assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
+      } else if (isMultiplexedSessionsEnabled(connection.getSpanner())) {
+        assertEquals(3, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
+      } else {
+        assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
+      }
     }
     assertEquals(2, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
     assertEquals(2, mockSpanner.countRequestsOfType(PartitionQueryRequest.class));