Skip to content

Commit 189ecbf

Browse files
harshachintaharshachinta
committed
chore(spanner): fix unit tests due to backgroung BeginTransaction RPC
1 parent dc12b5c commit 189ecbf

File tree

5 files changed

+38
-22
lines changed

5 files changed

+38
-22
lines changed

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -247,9 +247,12 @@ public void onSessionReady(SessionImpl session) {
247247
// only start the maintainer if we actually managed to create a session in the first
248248
// place.
249249
maintainer.start();
250+
250251
// initiate a begin transaction request to verify if read-write transactions are
251252
// supported using multiplexed sessions.
252-
verifyBeginTransactionWithRWOnMultiplexedSession(session.getName());
253+
if (sessionClient.getSpanner().getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()) {
254+
verifyBeginTransactionWithRWOnMultiplexedSession(session.getName());
255+
}
253256
}
254257

255258
@Override

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbstractMockServerTest.java‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -321,4 +321,17 @@ boolean isMultiplexedSessionsEnabled(Spanner spanner) {
321321
}
322322
return spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession();
323323
}
324+
325+
boolean isMultiplexedSessionsEnabledForRW(Spanner spanner) {
326+
if (spanner.getOptions() == null || spanner.getOptions().getSessionPoolOptions() == null) {
327+
return false;
328+
}
329+
return spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW();
330+
}
331+
332+
// Increment the total request count by 1 when multiplexed session is enabled for read-write transactions..
333+
// This is due to the explicit BeginTransaction RPC that is executed once during multiplexed session creation.
334+
protected int mayBeIncrementBeginTransactionRequestsCount(Spanner spanner, int count) {
335+
return isMultiplexedSessionsEnabledForRW(spanner) ? count + 1 : count;
336+
}
324337
}

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/DelayTransactionStartUntilFirstWriteMockServerTest.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -362,7 +362,7 @@ public void testTransactionWithQueryFollowedByMutations() {
362362
ExecuteSqlRequest queryRequest =
363363
mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0);
364364
assertFalse(queryRequest.hasTransaction());
365-
assertEquals(1, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
365+
assertEquals(mayBeIncrementBeginTransactionRequestsCount(connection.getSpanner(), /* count = */ 1), mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
366366
assertEquals(1, mockSpanner.countRequestsOfType(CommitRequest.class));
367367
}
368368
}
@@ -382,7 +382,7 @@ public void testTransactionWithMutationsFollowedByQuery() {
382382
ExecuteSqlRequest queryRequest =
383383
mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0);
384384
assertFalse(queryRequest.hasTransaction());
385-
assertEquals(1, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
385+
assertEquals(mayBeIncrementBeginTransactionRequestsCount(connection.getSpanner(), /* count = */ 1), mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
386386
assertEquals(1, mockSpanner.countRequestsOfType(CommitRequest.class));
387387
}
388388
}

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

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -98,8 +98,8 @@ public void testPartitionQuery() {
9898
} else {
9999
assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
100100
}
101+
assertEquals(mayBeIncrementBeginTransactionRequestsCount(connection.getSpanner(), /* count = */ 2), mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
101102
}
102-
assertEquals(2, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
103103
assertEquals(2, mockSpanner.countRequestsOfType(PartitionQueryRequest.class));
104104
}
105105

@@ -125,8 +125,8 @@ public void testPartitionQueryInReadOnlyTransaction() {
125125
} else {
126126
assertEquals(1, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
127127
}
128+
assertEquals(mayBeIncrementBeginTransactionRequestsCount(connection.getSpanner(), /* count = */ 1), mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
128129
}
129-
assertEquals(1, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
130130
assertEquals(2, mockSpanner.countRequestsOfType(PartitionQueryRequest.class));
131131
}
132132

@@ -160,12 +160,12 @@ public void testMixNormalAndPartitionQueryInReadOnlyTransaction() {
160160
} else {
161161
assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
162162
}
163+
assertEquals(mayBeIncrementBeginTransactionRequestsCount(connection.getSpanner(), /* count = */ 4), mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
163164
}
164165
// The above will start two transactions:
165166
// 1. The initial 'normal' read-only transaction.
166167
// 2. The batch read-only transaction. The latter will use the same read timestamp as the normal
167168
// read-only transaction.
168-
assertEquals(4, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
169169
assertEquals(2, mockSpanner.countRequestsOfType(PartitionQueryRequest.class));
170170

171171
List<BeginTransactionRequest> beginRequests =
@@ -233,7 +233,7 @@ public void testRunPartition() {
233233
expectedCreateSessionsRPC,
234234
mockSpanner.countRequestsOfType(CreateSessionRequest.class));
235235
}
236-
assertEquals(1, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
236+
assertEquals(mayBeIncrementBeginTransactionRequestsCount(connection.getSpanner(), /* count = */ 1), mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
237237
assertEquals(1, mockSpanner.countRequestsOfType(PartitionQueryRequest.class));
238238
assertEquals(maxPartitions, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
239239
assertFalse(
@@ -333,7 +333,7 @@ public void testRunPartitionUsingSql() {
333333
expectedCreateSessionsRPC,
334334
mockSpanner.countRequestsOfType(CreateSessionRequest.class));
335335
}
336-
assertEquals(1, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
336+
assertEquals(mayBeIncrementBeginTransactionRequestsCount(connection.getSpanner(), /* count = */ 1), mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
337337
assertEquals(1, mockSpanner.countRequestsOfType(PartitionQueryRequest.class));
338338
assertEquals(maxPartitions, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
339339
assertFalse(
@@ -379,8 +379,8 @@ public void testRunPartitionedQuery() {
379379
} else {
380380
assertEquals(1, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
381381
}
382+
assertEquals(mayBeIncrementBeginTransactionRequestsCount(connection.getSpanner(), /* count = */ 1), mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
382383
}
383-
assertEquals(1, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
384384
assertEquals(1, mockSpanner.countRequestsOfType(PartitionQueryRequest.class));
385385
}
386386

@@ -412,8 +412,8 @@ public void testRunEmptyPartitionedQuery() {
412412
} else {
413413
assertEquals(1, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
414414
}
415+
assertEquals(mayBeIncrementBeginTransactionRequestsCount(connection.getSpanner(), /* count = */ 1), mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
415416
}
416-
assertEquals(1, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
417417
assertEquals(1, mockSpanner.countRequestsOfType(PartitionQueryRequest.class));
418418
}
419419

@@ -447,8 +447,8 @@ public void testGetMetadataWithoutNextCall() {
447447

448448
assertFalse(resultSet.next());
449449
}
450+
assertEquals(mayBeIncrementBeginTransactionRequestsCount(connection.getSpanner(), /* count = */ 1), mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
450451
}
451-
assertEquals(1, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
452452
assertEquals(1, mockSpanner.countRequestsOfType(PartitionQueryRequest.class));
453453
}
454454

@@ -476,8 +476,8 @@ public void testGetMetadataWithoutNextCallOnEmptyResultSet() {
476476

477477
assertFalse(resultSet.next());
478478
}
479+
assertEquals(mayBeIncrementBeginTransactionRequestsCount(connection.getSpanner(), /* count = */ 1), mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
479480
}
480-
assertEquals(1, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
481481
assertEquals(1, mockSpanner.countRequestsOfType(PartitionQueryRequest.class));
482482
}
483483

@@ -500,8 +500,8 @@ public void testGetMetadataWithoutNextCallOnResultSetWithError() {
500500
assertThrows(SpannerException.class, resultSet::getMetadata);
501501
assertThrows(SpannerException.class, resultSet::getType);
502502
}
503+
assertEquals(mayBeIncrementBeginTransactionRequestsCount(connection.getSpanner(), /* count = */ 1), mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
503504
}
504-
assertEquals(1, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
505505
assertEquals(1, mockSpanner.countRequestsOfType(PartitionQueryRequest.class));
506506
}
507507

@@ -575,9 +575,9 @@ public void testRunPartitionedQueryUsingSql() {
575575
} else {
576576
assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
577577
}
578+
assertEquals(mayBeIncrementBeginTransactionRequestsCount(connection.getSpanner(), /* count = */ 2), mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
578579
}
579580
// We have 2 requests of each, as we run the query with data boost both enabled and disabled.
580-
assertEquals(2, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
581581
assertEquals(2, mockSpanner.countRequestsOfType(PartitionQueryRequest.class));
582582
}
583583

@@ -637,8 +637,8 @@ public void testRunPartitionedQueryWithError() {
637637
} else {
638638
assertEquals(1, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
639639
}
640+
assertEquals(mayBeIncrementBeginTransactionRequestsCount(connection.getSpanner(), /* count = */ 1), mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
640641
}
641-
assertEquals(1, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
642642
assertEquals(1, mockSpanner.countRequestsOfType(PartitionQueryRequest.class));
643643
}
644644

@@ -684,8 +684,8 @@ public void testRunPartitionedQueryWithMaxParallelism() {
684684
} else {
685685
assertEquals(5, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
686686
}
687+
assertEquals(mayBeIncrementBeginTransactionRequestsCount(connection.getSpanner(), /* count = */ 5), mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
687688
}
688-
assertEquals(5, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
689689
assertEquals(5, mockSpanner.countRequestsOfType(PartitionQueryRequest.class));
690690
}
691691

@@ -763,8 +763,8 @@ public void testAutoPartitionMode() {
763763
} else {
764764
assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class));
765765
}
766+
assertEquals(mayBeIncrementBeginTransactionRequestsCount(connection.getSpanner(), /* count = */ 2), mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
766767
}
767-
assertEquals(2, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
768768
assertEquals(2, mockSpanner.countRequestsOfType(PartitionQueryRequest.class));
769769
}
770770
}

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SavepointMockServerTest.java‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -231,7 +231,7 @@ public void testRollbackToSavepointInReadOnlyTransaction() {
231231
// It is a no-op on Cloud Spanner.
232232
connection.commit();
233233

234-
assertEquals(1, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
234+
assertEquals(mayBeIncrementBeginTransactionRequestsCount(connection.getSpanner(), 1), mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
235235
BeginTransactionRequest beginRequest =
236236
mockSpanner.getRequestsOfType(BeginTransactionRequest.class).get(0);
237237
assertTrue(beginRequest.getOptions().hasReadOnly());
@@ -266,7 +266,7 @@ public void testRollbackToSavepointInReadWriteTransaction() {
266266
connection.commit();
267267

268268
// Read/write transactions are started with inlined Begin transaction options.
269-
assertEquals(0, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
269+
assertEquals(mayBeIncrementBeginTransactionRequestsCount(connection.getSpanner(), 0), mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
270270
assertEquals(1, mockSpanner.countRequestsOfType(RollbackRequest.class));
271271
assertEquals(1, mockSpanner.countRequestsOfType(CommitRequest.class));
272272
assertEquals(2, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
@@ -312,7 +312,7 @@ public void testRollbackToSavepointWithDmlStatements() {
312312
connection.commit();
313313

314314
// Read/write transactions are started with inlined Begin transaction options.
315-
assertEquals(0, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
315+
assertEquals(mayBeIncrementBeginTransactionRequestsCount(connection.getSpanner(), 0), mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
316316
assertEquals(1, mockSpanner.countRequestsOfType(RollbackRequest.class));
317317
assertEquals(1, mockSpanner.countRequestsOfType(CommitRequest.class));
318318
assertEquals(5, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
@@ -368,7 +368,7 @@ public void testRollbackToSavepointFails() {
368368
assertThrows(AbortedDueToConcurrentModificationException.class, connection::commit);
369369

370370
// Read/write transactions are started with inlined Begin transaction options.
371-
assertEquals(0, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
371+
assertEquals(mayBeIncrementBeginTransactionRequestsCount(connection.getSpanner(), 0), mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
372372
assertEquals(2, mockSpanner.countRequestsOfType(RollbackRequest.class));
373373
assertEquals(0, mockSpanner.countRequestsOfType(CommitRequest.class));
374374
assertEquals(4, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
@@ -456,7 +456,7 @@ public void testRollbackToSavepointSucceedsWithRollback() {
456456
connection.rollback();
457457

458458
// Read/write transactions are started with inlined Begin transaction options.
459-
assertEquals(0, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
459+
assertEquals(mayBeIncrementBeginTransactionRequestsCount(connection.getSpanner(), 0), mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
460460
assertEquals(1, mockSpanner.countRequestsOfType(RollbackRequest.class));
461461
assertEquals(0, mockSpanner.countRequestsOfType(CommitRequest.class));
462462
assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));

0 commit comments

Comments
 (0)