diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/InlineBeginTransactionTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/InlineBeginTransactionTest.java
index a70e2a7aa34..65f4c5007ff 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/InlineBeginTransactionTest.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/InlineBeginTransactionTest.java
@@ -1668,192 +1668,203 @@ public void testWaitForTransactionTimeoutForCommit() {
       assertEquals(0, countRequests(CommitRequest.class));
     }
 
+    static void runWithIgnoreInlineBegin(Runnable runnable) {
+      // This will cause statements that requests a transaction to not return a transaction id.
+      mockSpanner.setIgnoreInlineBeginRequest(true);
+      try {
+        runnable.run();
+      } finally {
+        mockSpanner.setIgnoreInlineBeginRequest(false);
+      }
+    }
+
     @Test
     public void testQueryWithInlineBeginDidNotReturnTransaction() {
-      DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
-      // This will cause the first statement that requests a transaction to not return a transaction
-      // id.
-      mockSpanner.ignoreNextInlineBeginRequest();
-      SpannerException e =
-          assertThrows(
-              SpannerException.class,
-              () ->
-                  client
-                      .readWriteTransaction()
-                      .run(
-                          transaction -> {
-                            try (ResultSet rs =
-                                transaction.executeQuery(SELECT1_UNION_ALL_SELECT2)) {
-                              while (rs.next()) {}
-                            }
-                            return null;
-                          }));
-      assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode());
-      assertThat(e.getMessage()).contains(AbstractReadContext.NO_TRANSACTION_RETURNED_MSG);
-      assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(0);
-      assertThat(countRequests(ExecuteSqlRequest.class)).isEqualTo(1);
-      assertThat(countRequests(CommitRequest.class)).isEqualTo(0);
+      runWithIgnoreInlineBegin(
+          () -> {
+            DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+            SpannerException e =
+                assertThrows(
+                    SpannerException.class,
+                    () ->
+                        client
+                            .readWriteTransaction()
+                            .run(
+                                transaction -> {
+                                  try (ResultSet rs =
+                                      transaction.executeQuery(SELECT1_UNION_ALL_SELECT2)) {
+                                    while (rs.next()) {}
+                                  }
+                                  return null;
+                                }));
+            assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode());
+            assertThat(e.getMessage()).contains(AbstractReadContext.NO_TRANSACTION_RETURNED_MSG);
+            assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(0);
+            assertThat(countRequests(ExecuteSqlRequest.class)).isEqualTo(1);
+            assertThat(countRequests(CommitRequest.class)).isEqualTo(0);
+          });
     }
 
     @Test
     public void testReadWithInlineBeginDidNotReturnTransaction() {
-      DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
-      // This will cause the first statement that requests a transaction to not return a transaction
-      // id.
-      mockSpanner.ignoreNextInlineBeginRequest();
-      SpannerException e =
-          assertThrows(
-              SpannerException.class,
-              () ->
-                  client
-                      .readWriteTransaction()
-                      .run(
-                          transaction ->
-                              transaction.readRow(
-                                  "FOO", Key.of(1L), Collections.singletonList("BAR"))));
-      assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode());
-      assertThat(e.getMessage()).contains(AbstractReadContext.NO_TRANSACTION_RETURNED_MSG);
-      assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(0);
-      assertThat(countRequests(ReadRequest.class)).isEqualTo(1);
-      assertThat(countRequests(CommitRequest.class)).isEqualTo(0);
+      runWithIgnoreInlineBegin(
+          () -> {
+            DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+            SpannerException e =
+                assertThrows(
+                    SpannerException.class,
+                    () ->
+                        client
+                            .readWriteTransaction()
+                            .run(
+                                transaction ->
+                                    transaction.readRow(
+                                        "FOO", Key.of(1L), Collections.singletonList("BAR"))));
+            assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode());
+            assertThat(e.getMessage()).contains(AbstractReadContext.NO_TRANSACTION_RETURNED_MSG);
+            assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(0);
+            assertThat(countRequests(ReadRequest.class)).isEqualTo(1);
+            assertThat(countRequests(CommitRequest.class)).isEqualTo(0);
+          });
     }
 
     @Test
     public void testUpdateWithInlineBeginDidNotReturnTransaction() {
-      DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
-      // This will cause the first statement that requests a transaction to not return a transaction
-      // id.
-      mockSpanner.ignoreNextInlineBeginRequest();
-      SpannerException e =
-          assertThrows(
-              SpannerException.class,
-              () ->
-                  client
-                      .readWriteTransaction()
-                      .run(transaction -> transaction.executeUpdate(UPDATE_STATEMENT)));
-      assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode());
-      assertThat(e.getMessage()).contains(AbstractReadContext.NO_TRANSACTION_RETURNED_MSG);
-      assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(0);
-      assertThat(countRequests(ExecuteSqlRequest.class)).isEqualTo(1);
-      assertThat(countRequests(CommitRequest.class)).isEqualTo(0);
+      runWithIgnoreInlineBegin(
+          () -> {
+            DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+            SpannerException e =
+                assertThrows(
+                    SpannerException.class,
+                    () ->
+                        client
+                            .readWriteTransaction()
+                            .run(transaction -> transaction.executeUpdate(UPDATE_STATEMENT)));
+            assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode());
+            assertThat(e.getMessage()).contains(AbstractReadContext.NO_TRANSACTION_RETURNED_MSG);
+            assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(0);
+            assertThat(countRequests(ExecuteSqlRequest.class)).isEqualTo(1);
+            assertThat(countRequests(CommitRequest.class)).isEqualTo(0);
+          });
     }
 
     @Test
     public void testBatchUpdateWithInlineBeginDidNotReturnTransaction() {
-      DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
-      // This will cause the first statement that requests a transaction to not return a transaction
-      // id.
-      mockSpanner.ignoreNextInlineBeginRequest();
-      SpannerException e =
-          assertThrows(
-              SpannerException.class,
-              () ->
-                  client
-                      .readWriteTransaction()
-                      .run(
-                          transaction -> {
-                            transaction.batchUpdate(Collections.singletonList(UPDATE_STATEMENT));
-                            return null;
-                          }));
-      assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode());
-      assertThat(e.getMessage()).contains(AbstractReadContext.NO_TRANSACTION_RETURNED_MSG);
-      assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(0);
-      assertThat(countRequests(ExecuteBatchDmlRequest.class)).isEqualTo(1);
-      assertThat(countRequests(CommitRequest.class)).isEqualTo(0);
+      runWithIgnoreInlineBegin(
+          () -> {
+            DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+            SpannerException e =
+                assertThrows(
+                    SpannerException.class,
+                    () ->
+                        client
+                            .readWriteTransaction()
+                            .run(
+                                transaction -> {
+                                  transaction.batchUpdate(
+                                      Collections.singletonList(UPDATE_STATEMENT));
+                                  return null;
+                                }));
+            assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode());
+            assertThat(e.getMessage()).contains(AbstractReadContext.NO_TRANSACTION_RETURNED_MSG);
+            assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(0);
+            assertThat(countRequests(ExecuteBatchDmlRequest.class)).isEqualTo(1);
+            assertThat(countRequests(CommitRequest.class)).isEqualTo(0);
+          });
     }
 
     @Test
     public void testQueryAsyncWithInlineBeginDidNotReturnTransaction() {
-      DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
-      final ExecutorService executor = Executors.newSingleThreadExecutor();
-      // This will cause the first statement that requests a transaction to not return a transaction
-      // id.
-      mockSpanner.ignoreNextInlineBeginRequest();
-      SpannerException outerException =
-          assertThrows(
-              SpannerException.class,
-              () ->
-                  client
-                      .readWriteTransaction()
-                      .run(
-                          transaction -> {
-                            try (AsyncResultSet rs =
-                                transaction.executeQueryAsync(SELECT1_UNION_ALL_SELECT2)) {
-                              return SpannerApiFutures.get(
-                                  rs.setCallback(
-                                      executor,
-                                      resultSet -> {
-                                        try {
-                                          while (true) {
-                                            switch (resultSet.tryNext()) {
-                                              case OK:
-                                                break;
-                                              case DONE:
+      runWithIgnoreInlineBegin(
+          () -> {
+            DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+            final ExecutorService executor = Executors.newSingleThreadExecutor();
+            SpannerException outerException =
+                assertThrows(
+                    SpannerException.class,
+                    () ->
+                        client
+                            .readWriteTransaction()
+                            .run(
+                                transaction -> {
+                                  try (AsyncResultSet rs =
+                                      transaction.executeQueryAsync(SELECT1_UNION_ALL_SELECT2)) {
+                                    return SpannerApiFutures.get(
+                                        rs.setCallback(
+                                            executor,
+                                            resultSet -> {
+                                              try {
+                                                while (true) {
+                                                  switch (resultSet.tryNext()) {
+                                                    case OK:
+                                                      break;
+                                                    case DONE:
+                                                      return CallbackResponse.DONE;
+                                                    case NOT_READY:
+                                                      return CallbackResponse.CONTINUE;
+                                                  }
+                                                }
+                                              } catch (SpannerException e) {
                                                 return CallbackResponse.DONE;
-                                              case NOT_READY:
-                                                return CallbackResponse.CONTINUE;
-                                            }
-                                          }
-                                        } catch (SpannerException e) {
-                                          return CallbackResponse.DONE;
-                                        }
-                                      }));
-                            }
-                          }));
-      assertEquals(ErrorCode.FAILED_PRECONDITION, outerException.getErrorCode());
-      assertThat(outerException.getMessage())
-          .contains(AbstractReadContext.NO_TRANSACTION_RETURNED_MSG);
-
-      assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(0);
-      assertThat(countRequests(ExecuteSqlRequest.class)).isEqualTo(1);
-      assertThat(countRequests(CommitRequest.class)).isEqualTo(0);
+                                              }
+                                            }));
+                                  }
+                                }));
+            assertEquals(ErrorCode.FAILED_PRECONDITION, outerException.getErrorCode());
+            assertThat(outerException.getMessage())
+                .contains(AbstractReadContext.NO_TRANSACTION_RETURNED_MSG);
+
+            assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(0);
+            assertThat(countRequests(ExecuteSqlRequest.class)).isEqualTo(1);
+            assertThat(countRequests(CommitRequest.class)).isEqualTo(0);
+          });
     }
 
     @Test
     public void testUpdateAsyncWithInlineBeginDidNotReturnTransaction() {
-      DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
-      // This will cause the first statement that requests a transaction to not return a transaction
-      // id.
-      mockSpanner.ignoreNextInlineBeginRequest();
-      SpannerException e =
-          assertThrows(
-              SpannerException.class,
-              () ->
-                  client
-                      .readWriteTransaction()
-                      .run(
-                          transaction ->
-                              SpannerApiFutures.get(
-                                  transaction.executeUpdateAsync(UPDATE_STATEMENT))));
-      assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode());
-      assertThat(e.getMessage()).contains(AbstractReadContext.NO_TRANSACTION_RETURNED_MSG);
-      assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(0);
-      assertThat(countRequests(ExecuteSqlRequest.class)).isEqualTo(1);
-      assertThat(countRequests(CommitRequest.class)).isEqualTo(0);
+      runWithIgnoreInlineBegin(
+          () -> {
+            DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+            SpannerException e =
+                assertThrows(
+                    SpannerException.class,
+                    () ->
+                        client
+                            .readWriteTransaction()
+                            .run(
+                                transaction ->
+                                    SpannerApiFutures.get(
+                                        transaction.executeUpdateAsync(UPDATE_STATEMENT))));
+            assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode());
+            assertThat(e.getMessage()).contains(AbstractReadContext.NO_TRANSACTION_RETURNED_MSG);
+            assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(0);
+            assertThat(countRequests(ExecuteSqlRequest.class)).isEqualTo(1);
+            assertThat(countRequests(CommitRequest.class)).isEqualTo(0);
+          });
     }
 
     @Test
     public void testBatchUpdateAsyncWithInlineBeginDidNotReturnTransaction() {
-      DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
-      // This will cause the first statement that requests a transaction to not return a transaction
-      // id.
-      mockSpanner.ignoreNextInlineBeginRequest();
-      SpannerException e =
-          assertThrows(
-              SpannerException.class,
-              () ->
-                  client
-                      .readWriteTransaction()
-                      .run(
-                          transaction ->
-                              SpannerApiFutures.get(
-                                  transaction.batchUpdateAsync(
-                                      Collections.singletonList(UPDATE_STATEMENT)))));
-      assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode());
-      assertThat(e.getMessage()).contains(AbstractReadContext.NO_TRANSACTION_RETURNED_MSG);
-      assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(0);
-      assertThat(countRequests(ExecuteBatchDmlRequest.class)).isEqualTo(1);
-      assertThat(countRequests(CommitRequest.class)).isEqualTo(0);
+      runWithIgnoreInlineBegin(
+          () -> {
+            DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+            SpannerException e =
+                assertThrows(
+                    SpannerException.class,
+                    () ->
+                        client
+                            .readWriteTransaction()
+                            .run(
+                                transaction ->
+                                    SpannerApiFutures.get(
+                                        transaction.batchUpdateAsync(
+                                            Collections.singletonList(UPDATE_STATEMENT)))));
+            assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode());
+            assertThat(e.getMessage()).contains(AbstractReadContext.NO_TRANSACTION_RETURNED_MSG);
+            assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(0);
+            assertThat(countRequests(ExecuteBatchDmlRequest.class)).isEqualTo(1);
+            assertThat(countRequests(CommitRequest.class)).isEqualTo(0);
+          });
     }
 
     @Test
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerServiceImpl.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerServiceImpl.java
index ef637db594b..23280745592 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerServiceImpl.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerServiceImpl.java
@@ -607,7 +607,7 @@ private static void checkStreamException(
   private ConcurrentMap<ByteString, Boolean> commitRetryTransactions = new ConcurrentHashMap<>();
   private final AtomicBoolean abortNextTransaction = new AtomicBoolean();
   private final AtomicBoolean abortNextStatement = new AtomicBoolean();
-  private final AtomicBoolean ignoreNextInlineBeginRequest = new AtomicBoolean();
+  private final AtomicBoolean ignoreInlineBeginRequest = new AtomicBoolean();
   private ConcurrentMap<String, AtomicLong> transactionCounters = new ConcurrentHashMap<>();
   private ConcurrentMap<String, List<ByteString>> partitionTokens = new ConcurrentHashMap<>();
   private ConcurrentMap<ByteString, Instant> transactionLastUsed = new ConcurrentHashMap<>();
@@ -790,8 +790,8 @@ public void abortAllTransactions() {
     }
   }
 
-  public void ignoreNextInlineBeginRequest() {
-    ignoreNextInlineBeginRequest.set(true);
+  public void setIgnoreInlineBeginRequest(boolean ignore) {
+    ignoreInlineBeginRequest.set(ignore);
   }
 
   public void freeze() {
@@ -1063,7 +1063,7 @@ public void executeSql(ExecuteSqlRequest request, StreamObserver<ResultSet> resp
                     .setMetadata(
                         ResultSetMetadata.newBuilder()
                             .setTransaction(
-                                ignoreNextInlineBeginRequest.getAndSet(false)
+                                ignoreInlineBeginRequest.get()
                                     ? Transaction.getDefaultInstance()
                                     : Transaction.newBuilder().setId(transactionId).build())
                             .build());
@@ -1096,7 +1096,7 @@ private void returnResultSet(
           metadata
               .toBuilder()
               .setTransaction(
-                  ignoreNextInlineBeginRequest.getAndSet(false)
+                  ignoreInlineBeginRequest.get()
                       ? Transaction.getDefaultInstance()
                       : Transaction.newBuilder().setId(transactionId).build())
               .build();
@@ -1197,7 +1197,7 @@ public void executeBatchDml(
                 .setMetadata(
                     ResultSetMetadata.newBuilder()
                         .setTransaction(
-                            ignoreNextInlineBeginRequest.getAndSet(false)
+                            ignoreInlineBeginRequest.get()
                                 ? Transaction.getDefaultInstance()
                                 : Transaction.newBuilder().setId(transactionId).build())
                         .build())
@@ -1730,7 +1730,7 @@ private void returnPartialResultSet(
           metadata
               .toBuilder()
               .setTransaction(
-                  ignoreNextInlineBeginRequest.getAndSet(false)
+                  ignoreInlineBeginRequest.get()
                       ? Transaction.getDefaultInstance()
                       : Transaction.newBuilder().setId(transactionId).build())
               .build();
@@ -1785,7 +1785,7 @@ private void returnPartialResultSet(
                   ResultSetMetadata.newBuilder()
                       .setRowType(StructType.newBuilder().addFields(field).build())
                       .setTransaction(
-                          ignoreNextInlineBeginRequest.getAndSet(false)
+                          ignoreInlineBeginRequest.get()
                               ? Transaction.getDefaultInstance()
                               : Transaction.newBuilder().setId(transaction.getId()).build())
                       .build())
@@ -1798,7 +1798,7 @@ private void returnPartialResultSet(
                   ResultSetMetadata.newBuilder()
                       .setRowType(StructType.newBuilder().addFields(field).build())
                       .setTransaction(
-                          ignoreNextInlineBeginRequest.getAndSet(false)
+                          ignoreInlineBeginRequest.get()
                               ? Transaction.getDefaultInstance()
                               : Transaction.newBuilder().setId(transaction.getId()).build())
                       .build())