chore(spanner): add unit tests in TransactionManagerImplTest

Author: harshachinta

google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java

@@ -209,7 +209,7 @@ public void removeListener(Runnable listener) {
 
     volatile ByteString transactionId;
 
-    final ByteString previousAbortedTransactionId;
+    ByteString previousAbortedTransactionId;
 
     private CommitResponse commitResponse;
     private final Clock clock;

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

@@ -98,7 +98,7 @@ public void setUp() {
 
   @Test
   public void beginCalledTwiceFails() {
-    when(session.newTransaction(Options.fromTransactionOptions())).thenReturn(txn);
+    when(session.newTransaction(Options.fromTransactionOptions(), null)).thenReturn(txn);
     assertThat(manager.begin()).isEqualTo(txn);
     assertThat(manager.getState()).isEqualTo(TransactionState.STARTED);
     IllegalStateException e = assertThrows(IllegalStateException.class, () -> manager.begin());
@@ -126,7 +126,7 @@ public void resetBeforeBeginFails() {
 
   @Test
   public void transactionRolledBackOnClose() {
-    when(session.newTransaction(Options.fromTransactionOptions())).thenReturn(txn);
+    when(session.newTransaction(Options.fromTransactionOptions(), null)).thenReturn(txn);
     when(txn.isAborted()).thenReturn(false);
     manager.begin();
     manager.close();
@@ -135,7 +135,7 @@ public void transactionRolledBackOnClose() {
 
   @Test
   public void commitSucceeds() {
-    when(session.newTransaction(Options.fromTransactionOptions())).thenReturn(txn);
+    when(session.newTransaction(Options.fromTransactionOptions(), null)).thenReturn(txn);
     Timestamp commitTimestamp = Timestamp.ofTimeMicroseconds(1);
     CommitResponse response = new CommitResponse(commitTimestamp);
     when(txn.getCommitResponse()).thenReturn(response);
@@ -147,7 +147,7 @@ public void commitSucceeds() {
 
   @Test
   public void resetAfterSuccessfulCommitFails() {
-    when(session.newTransaction(Options.fromTransactionOptions())).thenReturn(txn);
+    when(session.newTransaction(Options.fromTransactionOptions(), null)).thenReturn(txn);
     manager.begin();
     manager.commit();
     IllegalStateException e =
@@ -157,21 +157,21 @@ public void resetAfterSuccessfulCommitFails() {
 
   @Test
   public void resetAfterAbortSucceeds() {
-    when(session.newTransaction(Options.fromTransactionOptions())).thenReturn(txn);
+    when(session.newTransaction(Options.fromTransactionOptions(), null)).thenReturn(txn);
     manager.begin();
     doThrow(SpannerExceptionFactory.newSpannerException(ErrorCode.ABORTED, "")).when(txn).commit();
     assertThrows(AbortedException.class, () -> manager.commit());
     assertEquals(TransactionState.ABORTED, manager.getState());
 
     txn = Mockito.mock(TransactionRunnerImpl.TransactionContextImpl.class);
-    when(session.newTransaction(Options.fromTransactionOptions())).thenReturn(txn);
+    when(session.newTransaction(Options.fromTransactionOptions(), null)).thenReturn(txn);
     assertThat(manager.resetForRetry()).isEqualTo(txn);
     assertThat(manager.getState()).isEqualTo(TransactionState.STARTED);
   }
 
   @Test
   public void resetAfterErrorFails() {
-    when(session.newTransaction(Options.fromTransactionOptions())).thenReturn(txn);
+    when(session.newTransaction(Options.fromTransactionOptions(), null)).thenReturn(txn);
     manager.begin();
     doThrow(SpannerExceptionFactory.newSpannerException(ErrorCode.UNKNOWN, "")).when(txn).commit();
     SpannerException e = assertThrows(SpannerException.class, () -> manager.commit());
@@ -184,7 +184,7 @@ public void resetAfterErrorFails() {
 
   @Test
   public void rollbackAfterCommitFails() {
-    when(session.newTransaction(Options.fromTransactionOptions())).thenReturn(txn);
+    when(session.newTransaction(Options.fromTransactionOptions(), null)).thenReturn(txn);
     manager.begin();
     manager.commit();
     IllegalStateException e = assertThrows(IllegalStateException.class, () -> manager.rollback());
@@ -193,7 +193,7 @@ public void rollbackAfterCommitFails() {
 
   @Test
   public void commitAfterRollbackFails() {
-    when(session.newTransaction(Options.fromTransactionOptions())).thenReturn(txn);
+    when(session.newTransaction(Options.fromTransactionOptions(), null)).thenReturn(txn);
     manager.begin();
     manager.rollback();
     IllegalStateException e = assertThrows(IllegalStateException.class, () -> manager.commit());
@@ -363,4 +363,42 @@ public void inlineBegin() {
       assertThat(transactionsStarted.get()).isEqualTo(1);
     }
   }
+
+  @Test
+  public void storePreviousTxnIdOnAbortForMultiplexedSession() {
+    txn = Mockito.mock(TransactionRunnerImpl.TransactionContextImpl.class);
+    final ByteString mockTransactionId = ByteString.copyFromUtf8("mockTransactionId");
+    txn.transactionId = mockTransactionId;
+    when(session.newTransaction(Options.fromTransactionOptions(), null)).thenReturn(txn);
+    manager.begin();
+    doThrow(SpannerExceptionFactory.newSpannerException(ErrorCode.ABORTED, "")).when(txn).commit();
+    assertThrows(AbortedException.class, () -> manager.commit());
+    assertEquals(TransactionState.ABORTED, manager.getState());
+
+    txn = Mockito.mock(TransactionRunnerImpl.TransactionContextImpl.class);
+    txn.previousAbortedTransactionId = mockTransactionId;
+    when(session.newTransaction(Options.fromTransactionOptions(), mockTransactionId))
+        .thenReturn(txn);
+    when(session.getIsMultiplexed()).thenReturn(true);
+    assertThat(manager.resetForRetry()).isEqualTo(txn);
+    assertThat(manager.getState()).isEqualTo(TransactionState.STARTED);
+  }
+
+  @Test
+  public void skipTxnIdStorageOnAbortForRegularSession() {
+    txn = Mockito.mock(TransactionRunnerImpl.TransactionContextImpl.class);
+    final ByteString mockTransactionId = ByteString.copyFromUtf8("mockTransactionId");
+    txn.transactionId = mockTransactionId;
+    when(session.newTransaction(Options.fromTransactionOptions(), null)).thenReturn(txn);
+    manager.begin();
+    doThrow(SpannerExceptionFactory.newSpannerException(ErrorCode.ABORTED, "")).when(txn).commit();
+    assertThrows(AbortedException.class, () -> manager.commit());
+    assertEquals(TransactionState.ABORTED, manager.getState());
+
+    txn = Mockito.mock(TransactionRunnerImpl.TransactionContextImpl.class);
+    when(session.newTransaction(Options.fromTransactionOptions(), null)).thenReturn(txn);
+    when(session.getIsMultiplexed()).thenReturn(false);
+    assertThat(manager.resetForRetry()).isEqualTo(txn);
+    assertThat(manager.getState()).isEqualTo(TransactionState.STARTED);
+  }
 }