From 7564c11dcf4e06a85d288c88754fe0f0acf47e26 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= <koloite@gmail.com>
Date: Mon, 3 May 2021 23:53:45 +0200
Subject: [PATCH] test: move begin() outside retry loop (#1134)

The begin() call should be before the retry loop, and not inside. The
catch clause for an AbortedException should reset the transaction manager
and assign the new transaction context to the original transaction context
variable.

Fixes #1125
---
 .../cloud/spanner/DatabaseClientImplTest.java | 25 +++++++++++--------
 1 file changed, 14 insertions(+), 11 deletions(-)

diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java
index ed0af902a9..991e2c6cb6 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java
@@ -894,32 +894,34 @@ public void runAsyncWithException() throws Exception {
     executor.shutdown();
   }
 
+  @SuppressWarnings("resource")
   @Test
   public void testTransactionManager() {
     DatabaseClient client =
         spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
     try (TransactionManager manager = client.transactionManager()) {
+      TransactionContext transaction = manager.begin();
       while (true) {
-        TransactionContext transaction = manager.begin();
         try {
           transaction.executeUpdate(UPDATE_STATEMENT);
           manager.commit();
           assertNotNull(manager.getCommitTimestamp());
           break;
         } catch (AbortedException e) {
-          manager.resetForRetry();
+          transaction = manager.resetForRetry();
         }
       }
     }
   }
 
+  @SuppressWarnings("resource")
   @Test
   public void testTransactionManager_returnsCommitStats() {
     DatabaseClient client =
         spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
     try (TransactionManager manager = client.transactionManager(Options.commitStats())) {
+      TransactionContext transaction = manager.begin();
       while (true) {
-        TransactionContext transaction = manager.begin();
         try {
           transaction.buffer(Mutation.delete("FOO", Key.of("foo")));
           manager.commit();
@@ -928,12 +930,13 @@ public void testTransactionManager_returnsCommitStats() {
           assertEquals(1L, manager.getCommitResponse().getCommitStats().getMutationCount());
           break;
         } catch (AbortedException e) {
-          manager.resetForRetry();
+          transaction = manager.resetForRetry();
         }
       }
     }
   }
 
+  @SuppressWarnings("resource")
   @Test
   public void transactionManagerIsNonBlocking() throws Exception {
     mockSpanner.freeze();
@@ -941,16 +944,16 @@ public void transactionManagerIsNonBlocking() throws Exception {
         spannerWithEmptySessionPool.getDatabaseClient(
             DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
     try (TransactionManager txManager = client.transactionManager()) {
+      mockSpanner.unfreeze();
+      TransactionContext transaction = txManager.begin();
       while (true) {
-        mockSpanner.unfreeze();
-        TransactionContext tx = txManager.begin();
         try {
-          tx.executeUpdate(UPDATE_STATEMENT);
+          transaction.executeUpdate(UPDATE_STATEMENT);
           txManager.commit();
           break;
         } catch (AbortedException e) {
           Thread.sleep(e.getRetryDelayInMillis());
-          txManager.resetForRetry();
+          transaction = txManager.resetForRetry();
         }
       }
     }
@@ -962,10 +965,10 @@ public void transactionManagerExecuteQueryAsync() throws Exception {
         spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
     final AtomicInteger rowCount = new AtomicInteger();
     try (TransactionManager txManager = client.transactionManager()) {
+      TransactionContext transaction = txManager.begin();
       while (true) {
-        TransactionContext tx = txManager.begin();
         try {
-          try (AsyncResultSet rs = tx.executeQueryAsync(SELECT1)) {
+          try (AsyncResultSet rs = transaction.executeQueryAsync(SELECT1)) {
             rs.setCallback(
                 executor,
                 resultSet -> {
@@ -990,7 +993,7 @@ public void transactionManagerExecuteQueryAsync() throws Exception {
           break;
         } catch (AbortedException e) {
           Thread.sleep(e.getRetryDelayInMillis());
-          txManager.resetForRetry();
+          transaction = txManager.resetForRetry();
         }
       }
     }