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

Author: olavloite

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,29 +930,30 @@ 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();
     DatabaseClient client =
         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();
         }
       }
     }