Skip to content

Commit 136d31e

Browse files
harshachintaharshachinta
committed
chore(spanner): update test cases
1 parent 8cc2b32 commit 136d31e

File tree

3 files changed

+189
-5
lines changed

3 files changed

+189
-5
lines changed

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerImplTest.java‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,14 +16,18 @@
1616

1717
package com.google.cloud.spanner;
1818

19+
import static org.junit.Assert.assertThrows;
1920
import static org.mockito.ArgumentMatchers.any;
2021
import static org.mockito.ArgumentMatchers.eq;
22+
import static org.mockito.Mockito.clearInvocations;
23+
import static org.mockito.Mockito.doThrow;
2124
import static org.mockito.Mockito.mock;
2225
import static org.mockito.Mockito.verify;
2326
import static org.mockito.Mockito.when;
2427

2528
import com.google.api.core.ApiFutures;
2629
import com.google.cloud.Timestamp;
30+
import com.google.protobuf.ByteString;
2731
import io.opentelemetry.api.trace.Span;
2832
import io.opentelemetry.context.Scope;
2933
import org.junit.Test;
@@ -56,4 +60,65 @@ public void testCommitReturnsCommitStats() {
5660
verify(transaction).commitAsync();
5761
}
5862
}
63+
64+
@Test
65+
public void testRetryUsesPreviousTransactionIdOnMultiplexedSession() {
66+
// Set up mock transaction IDs
67+
final ByteString mockTransactionId = ByteString.copyFromUtf8("mockTransactionId");
68+
final ByteString mockPreviousTransactionId =
69+
ByteString.copyFromUtf8("mockPreviousTransactionId");
70+
71+
Span oTspan = mock(Span.class);
72+
ISpan span = new OpenTelemetrySpan(oTspan);
73+
when(oTspan.makeCurrent()).thenReturn(mock(Scope.class));
74+
// Mark the session as multiplexed.
75+
when(session.getIsMultiplexed()).thenReturn(true);
76+
77+
// Initialize a mock transaction with transactionId = null, previousTransactionId = null.
78+
transaction = mock(TransactionRunnerImpl.TransactionContextImpl.class);
79+
when(transaction.ensureTxnAsync()).thenReturn(ApiFutures.immediateFuture(null));
80+
when(session.newTransaction(eq(Options.fromTransactionOptions(Options.commitStats())), any()))
81+
.thenReturn(transaction);
82+
83+
// Simulate an ABORTED error being thrown when `commitAsync()` is called.
84+
doThrow(SpannerExceptionFactory.newSpannerException(ErrorCode.ABORTED, ""))
85+
.when(transaction)
86+
.commitAsync();
87+
88+
try (AsyncTransactionManagerImpl manager =
89+
new AsyncTransactionManagerImpl(session, span, Options.commitStats())) {
90+
manager.beginAsync();
91+
92+
// Verify that for the first transaction attempt, the `previousTransactionId` is null.
93+
// This is because no transaction has been previously aborted at this point.
94+
verify(session).newTransaction(Options.fromTransactionOptions(Options.commitStats()), null);
95+
assertThrows(AbortedException.class, manager::commitAsync);
96+
clearInvocations(session);
97+
98+
// Mock the transaction object to contain transactionID=null and
99+
// previousTransactionId=mockPreviousTransactionId
100+
transaction.previousTransactionId = mockPreviousTransactionId;
101+
manager.resetForRetryAsync();
102+
// Verify that in the first retry attempt, the `previousTransactionId` is passed to the new
103+
// transaction.
104+
// This allows Spanner to retry the transaction using the ID of the aborted transaction.
105+
verify(session)
106+
.newTransaction(
107+
Options.fromTransactionOptions(Options.commitStats()), mockPreviousTransactionId);
108+
assertThrows(AbortedException.class, manager::commitAsync);
109+
clearInvocations(session);
110+
111+
// Mock the transaction object to contain transactionID=mockTransactionId and
112+
// previousTransactionId=mockPreviousTransactionId and transactionID = null
113+
transaction.transactionId = mockTransactionId;
114+
manager.resetForRetryAsync();
115+
// Verify that the current `transactionId` is used in the retry.
116+
// This ensures the retry logic is working as expected with the latest transaction ID.
117+
verify(session)
118+
.newTransaction(Options.fromTransactionOptions(Options.commitStats()), mockTransactionId);
119+
120+
when(transaction.rollbackAsync()).thenReturn(ApiFutures.immediateFuture(null));
121+
manager.closeAsync();
122+
}
123+
}
59124
}

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientMockServerTest.java‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,9 @@
1616

1717
package com.google.cloud.spanner;
1818

19+
import static com.google.cloud.spanner.MockSpannerTestUtil.INVALID_UPDATE_STATEMENT;
20+
import static com.google.cloud.spanner.MockSpannerTestUtil.UPDATE_COUNT;
21+
import static com.google.cloud.spanner.MockSpannerTestUtil.UPDATE_STATEMENT;
1922
import static com.google.common.truth.Truth.assertThat;
2023
import static org.junit.Assert.assertEquals;
2124
import static org.junit.Assert.assertFalse;
@@ -58,6 +61,12 @@ public class MultiplexedSessionDatabaseClientMockServerTest extends AbstractMock
5861
public static void setupResults() {
5962
mockSpanner.putStatementResults(
6063
StatementResult.query(STATEMENT, new RandomResultSetGenerator(1).generate()));
64+
65+
mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT));
66+
mockSpanner.putStatementResult(
67+
StatementResult.exception(
68+
INVALID_UPDATE_STATEMENT,
69+
Status.INVALID_ARGUMENT.withDescription("invalid statement").asRuntimeException()));
6170
}
6271

6372
@Before
@@ -467,6 +476,106 @@ public void testWriteAtLeastOnceWithExcludeTxnFromChangeStreams() {
467476
assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get());
468477
}
469478

479+
// TODO(sriharshach): Uncomment test once Lock order preservation proto is published
480+
/*
481+
@Test
482+
public void testAbortedReadWriteTxnUsesPreviousTxnIdOnRetryWithInlineBegin() throws InterruptedException {
483+
DatabaseClientImpl client =
484+
(DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
485+
// Force the Commit RPC to return Aborted the first time it is called. The exception is cleared
486+
// after the first call, so the retry should succeed.
487+
mockSpanner.setCommitExecutionTime(
488+
SimulatedExecutionTime.ofException(
489+
mockSpanner.createAbortedException(ByteString.copyFromUtf8("test"))));
490+
Thread.sleep(10000);
491+
TransactionRunner runner = client.readWriteTransaction();
492+
runner.run(
493+
transaction -> {
494+
try (ResultSet resultSet =
495+
transaction.executeQuery(STATEMENT)) {
496+
while (resultSet.next()) {}
497+
}
498+
return null;
499+
});
500+
501+
List<ExecuteSqlRequest> executeSqlRequests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
502+
assertEquals(2, executeSqlRequests.size());
503+
504+
// Verify the requests are executed using multiplexed sessions
505+
for (ExecuteSqlRequest request : executeSqlRequests) {
506+
assertTrue(mockSpanner.getSession(request.getSession()).getMultiplexed());
507+
}
508+
509+
// Verify that the first request uses inline begin, and the previous transaction ID is set to ByteString.EMPTY
510+
assertTrue(executeSqlRequests.get(0).hasTransaction());
511+
assertTrue(executeSqlRequests.get(0).getTransaction().hasBegin());
512+
assertTrue(executeSqlRequests.get(0).getTransaction().getBegin().hasReadWrite());
513+
assertNotNull(executeSqlRequests.get(0).getTransaction().getBegin().getReadWrite()
514+
.getMultiplexedSessionPreviousTransactionId());
515+
assertThat(executeSqlRequests.get(0).getTransaction().getBegin().getReadWrite().getMultiplexedSessionPreviousTransactionId()).isEqualTo(ByteString.EMPTY);
516+
517+
// Verify that the second request uses inline begin, and the previous transaction ID is set appropriately
518+
assertTrue(executeSqlRequests.get(1).hasTransaction());
519+
assertTrue(executeSqlRequests.get(1).getTransaction().hasBegin());
520+
assertTrue(executeSqlRequests.get(1).getTransaction().getBegin().hasReadWrite());
521+
assertNotNull(executeSqlRequests.get(1).getTransaction().getBegin().getReadWrite()
522+
.getMultiplexedSessionPreviousTransactionId());
523+
assertThat(executeSqlRequests.get(1).getTransaction().getBegin().getReadWrite().getMultiplexedSessionPreviousTransactionId()).isNotEqualTo(ByteString.EMPTY);
524+
}
525+
*/
526+
527+
// TODO(sriharshach): Uncomment test once Lock order preservation proto is published
528+
/*
529+
@Test
530+
public void testAbortedReadWriteTxnUsesPreviousTxnIdOnRetryWithExplicitBegin() throws InterruptedException {
531+
DatabaseClientImpl client =
532+
(DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
533+
// Force the Commit RPC to return Aborted the first time it is called. The exception is cleared
534+
// after the first call, so the retry should succeed.
535+
mockSpanner.setCommitExecutionTime(
536+
SimulatedExecutionTime.ofException(
537+
mockSpanner.createAbortedException(ByteString.copyFromUtf8("test"))));
538+
Thread.sleep(10000);
539+
TransactionRunner runner = client.readWriteTransaction();
540+
Long updateCount = runner.run(
541+
transaction -> {
542+
// This update statement carries the BeginTransaction, but fails. This will
543+
// cause the entire transaction to be retried with an explicit
544+
// BeginTransaction RPC to ensure all statements in the transaction are
545+
// actually executed against the same transaction.
546+
SpannerException e =
547+
assertThrows(
548+
SpannerException.class,
549+
() -> transaction.executeUpdate(INVALID_UPDATE_STATEMENT));
550+
assertEquals(ErrorCode.INVALID_ARGUMENT, e.getErrorCode());
551+
return transaction.executeUpdate(UPDATE_STATEMENT);
552+
});
553+
554+
assertThat(updateCount).isEqualTo(1L);
555+
List<BeginTransactionRequest> beginTransactionRequests = mockSpanner.getRequestsOfType(BeginTransactionRequest.class);
556+
assertEquals(2, beginTransactionRequests.size());
557+
558+
// Verify the requests are executed using multiplexed sessions
559+
for (BeginTransactionRequest request : beginTransactionRequests) {
560+
assertTrue(mockSpanner.getSession(request.getSession()).getMultiplexed());
561+
}
562+
563+
// Verify that explicit begin transaction is called during retry, and the previous transaction ID is set to ByteString.EMPTY
564+
assertTrue(beginTransactionRequests.get(0).hasOptions());
565+
assertTrue(beginTransactionRequests.get(0).getOptions().hasReadWrite());
566+
assertNotNull(beginTransactionRequests.get(0).getOptions().getReadWrite()
567+
.getMultiplexedSessionPreviousTransactionId());
568+
assertThat(beginTransactionRequests.get(0).getOptions().getReadWrite().getMultiplexedSessionPreviousTransactionId()).isEqualTo(ByteString.EMPTY);
569+
570+
// The previous transaction with id (txn1) fails during commit operation with ABORTED error.
571+
// Verify that explicit begin transaction is called during retry, and the previous transaction ID is not ByteString.EMPTY (should be set to txn1)
572+
assertTrue(beginTransactionRequests.get(1).hasOptions());
573+
assertTrue(beginTransactionRequests.get(1).getOptions().hasReadWrite());
574+
assertNotNull(beginTransactionRequests.get(1).getOptions().getReadWrite().getMultiplexedSessionPreviousTransactionId());
575+
assertThat(beginTransactionRequests.get(1).getOptions().getReadWrite().getMultiplexedSessionPreviousTransactionId()).isNotEqualTo(ByteString.EMPTY);
576+
}
577+
*/
578+
470579
private void waitForSessionToBeReplaced(DatabaseClientImpl client) {
471580
assertNotNull(client.multiplexedSessionDatabaseClient);
472581
SessionReference sessionReference =

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java‎

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import static org.junit.Assert.assertThrows;
2323
import static org.mockito.ArgumentMatchers.any;
2424
import static org.mockito.ArgumentMatchers.eq;
25+
import static org.mockito.Mockito.clearInvocations;
2526
import static org.mockito.Mockito.doThrow;
2627
import static org.mockito.Mockito.eq;
2728
import static org.mockito.Mockito.mock;
@@ -375,17 +376,22 @@ public void storePreviousTxnIdOnAbortForMultiplexedSession() {
375376
txn.transactionId = mockTransactionId;
376377
when(session.newTransaction(Options.fromTransactionOptions(), null)).thenReturn(txn);
377378
manager.begin();
379+
// Verify that for the first transaction attempt, the `previousTransactionId` is null.
380+
// This is because no transaction has been previously aborted at this point.
381+
verify(session).newTransaction(Options.fromTransactionOptions(), null);
378382
doThrow(SpannerExceptionFactory.newSpannerException(ErrorCode.ABORTED, "")).when(txn).commit();
379383
assertThrows(AbortedException.class, () -> manager.commit());
380-
assertEquals(TransactionState.ABORTED, manager.getState());
381384

382385
txn = Mockito.mock(TransactionRunnerImpl.TransactionContextImpl.class);
383-
txn.previousAbortedTransactionId = mockTransactionId;
386+
txn.previousTransactionId = mockTransactionId;
384387
when(session.newTransaction(Options.fromTransactionOptions(), mockTransactionId))
385388
.thenReturn(txn);
386389
when(session.getIsMultiplexed()).thenReturn(true);
387390
assertThat(manager.resetForRetry()).isEqualTo(txn);
388-
assertThat(manager.getState()).isEqualTo(TransactionState.STARTED);
391+
// Verify that in the first retry attempt, the `previousTransactionId` is passed to the new
392+
// transaction.
393+
// This allows Spanner to retry the transaction using the ID of the aborted transaction.
394+
verify(session).newTransaction(Options.fromTransactionOptions(), mockTransactionId);
389395
}
390396

391397
// This test ensures that when a transaction is aborted in a regular session,
@@ -398,14 +404,18 @@ public void skipTxnIdStorageOnAbortForRegularSession() {
398404
txn.transactionId = mockTransactionId;
399405
when(session.newTransaction(Options.fromTransactionOptions(), null)).thenReturn(txn);
400406
manager.begin();
407+
verify(session).newTransaction(Options.fromTransactionOptions(), null);
401408
doThrow(SpannerExceptionFactory.newSpannerException(ErrorCode.ABORTED, "")).when(txn).commit();
402409
assertThrows(AbortedException.class, () -> manager.commit());
403-
assertEquals(TransactionState.ABORTED, manager.getState());
410+
clearInvocations(session);
404411

405412
txn = Mockito.mock(TransactionRunnerImpl.TransactionContextImpl.class);
406413
when(session.newTransaction(Options.fromTransactionOptions(), null)).thenReturn(txn);
407414
when(session.getIsMultiplexed()).thenReturn(false);
408415
assertThat(manager.resetForRetry()).isEqualTo(txn);
409-
assertThat(manager.getState()).isEqualTo(TransactionState.STARTED);
416+
// Verify that in the first retry attempt, the `previousTransactionId` is not passed to the new
417+
// transaction
418+
// in case of regular sessions.
419+
verify(session).newTransaction(Options.fromTransactionOptions(), null);
410420
}
411421
}

0 commit comments

Comments
 (0)