chore(spanner): fix code and tests

Author: harshachinta

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

@@ -55,7 +55,7 @@ public void setSpan(ISpan span) {
 
   @Override
   public void close() {
-    closeAsync();
+    SpannerApiFutures.get(closeAsync());
   }
 
   @Override
@@ -183,6 +183,10 @@ public ApiFuture<Void> rollbackAsync() {
 
   @Override
   public TransactionContextFuture resetForRetryAsync() {
+    if (txn == null || !txn.isAborted() && txnState != TransactionState.ABORTED) {
+      throw new IllegalStateException(
+          "resetForRetry can only be called if the previous attempt aborted");
+    }
     return new TransactionContextFutureImpl(this, internalBeginAsync(false));
   }
 

google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncExamplesTest.java

@@ -253,6 +253,12 @@ public void runAsync() throws Exception {
             },
             executor);
     assertThat(insertCount.get()).isEqualTo(1L);
+    if (env.getTestHelper().getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()) {
+      // The runAsync() method should only be called once on the runner.
+      // However, due to a bug in regular sessions, it can be executed multiple times on the same
+      // runner.
+      runner = client.runAsync();
+    }
     ApiFuture<Long> deleteCount =
         runner.runAsync(
             txn ->

google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITBatchDmlTest.java

@@ -17,6 +17,7 @@
 package com.google.cloud.spanner.it;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assume.assumeFalse;
 
 import com.google.api.gax.longrunning.OperationFuture;
 import com.google.cloud.spanner.Database;
@@ -84,6 +85,10 @@ public void dropTable() throws Exception {
 
   @Test
   public void noStatementsInRequest() {
+    // TODO(sriharshach): Remove this skip once backend support empty transactions to commit.
+    assumeFalse(
+        "Skipping for multiplexed sessions since it does not allow empty transactions to commit",
+        isUsingMultiplexedSessionsForRW());
     final TransactionCallable<long[]> callable =
         transaction -> {
           List<Statement> stmts = new ArrayList<>();
@@ -252,4 +257,8 @@ public void largeBatchDml_withNonParameterisedStatements() {
     assertThat(actualRowCounts.length).isEqualTo(80);
     assertThat(expectedRowCounts).isEqualTo(actualRowCounts);
   }
+
+  boolean isUsingMultiplexedSessionsForRW() {
+    return env.getTestHelper().getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW();
+  }
 }

google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITBatchReadTest.java

@@ -181,7 +181,9 @@ public static void setUpDatabase() throws Exception {
           totalSize = 0;
         }
       }
-      dbClient.write(mutations);
+      if (mutations.size() > 0) {
+        dbClient.write(mutations);
+      }
     }
     // Our read/queries are executed with some staleness.
     Thread.sleep(2 * STALENESS_MILLISEC);

google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITTransactionTest.java

@@ -464,7 +464,10 @@ public void nestedSingleUseReadTxnThrows() {
   @Test
   public void nestedTxnSucceedsWhenAllowed() {
     assumeFalse("Emulator does not support multiple parallel transactions", isUsingEmulator());
-
+    // TODO(sriharshach): Remove this skip once backend support empty transactions to commit.
+    assumeFalse(
+        "Skipping for multiplexed sessions since it does not allow empty transactions to commit",
+        isUsingMultiplexedSessionsForRW());
     client
         .readWriteTransaction()
         .allowNestedTransaction()
@@ -588,4 +591,8 @@ public void testTransactionRunnerReturnsCommitStats() {
     // MutationCount = 2 (2 columns).
     assertEquals(2L, runner.getCommitResponse().getCommitStats().getMutationCount());
   }
+
+  boolean isUsingMultiplexedSessionsForRW() {
+    return env.getTestHelper().getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW();
+  }
 }