test: add tests for error during tx

@@ -38,7 +38,10 @@
 import com.google.protobuf.AbstractMessage;
 import com.google.protobuf.ListValue;
 import com.google.spanner.v1.BeginTransactionRequest;
+import com.google.spanner.v1.CommitRequest;
+import com.google.spanner.v1.ExecuteSqlRequest;
 import com.google.spanner.v1.ResultSetMetadata;
+import com.google.spanner.v1.RollbackRequest;
 import com.google.spanner.v1.StructType;
 import com.google.spanner.v1.StructType.Field;
 import com.google.spanner.v1.TypeCode;
@@ -284,6 +287,63 @@ public Long run(TransactionContext transaction) throws Exception {
 assertThat(countTransactionsStarted()).isEqualTo(2);
 }
 
+@Test
+public void testInlinedBeginTxWithUncaughtError() {
+DatabaseClient client =
+spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]"));
+try {
+client
+.readWriteTransaction()
+.run(
+new TransactionCallable<Long>() {
+@Override
+public Long run(TransactionContext transaction) throws Exception {
+return transaction.executeUpdate(INVALID_UPDATE_STATEMENT);
+}
+});
+fail("missing expected exception");
+} catch (SpannerException e) {
+assertThat(e.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT);
+}
+// The first update will start a transaction, but then fail the update statement. This will
+// start a transaction on the mock server, but that transaction will never be returned to the
+// client.
+assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(0);
+assertThat(countRequests(CommitRequest.class)).isEqualTo(0);
+assertThat(countRequests(ExecuteSqlRequest.class)).isEqualTo(1);
+// No rollback request will be initiated because the client does not receive any transaction id.
+assertThat(countRequests(RollbackRequest.class)).isEqualTo(0);
+assertThat(countTransactionsStarted()).isEqualTo(1);
+}
+
+@Test
+public void testInlinedBeginTxWithUncaughtErrorAfterSuccessfulBegin() {
+DatabaseClient client =
+spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]"));
+try {
+client
+.readWriteTransaction()
+.run(
+new TransactionCallable<Long>() {
+@Override
+public Long run(TransactionContext transaction) throws Exception {
+// This statement will start a transaction.
+transaction.executeUpdate(UPDATE_STATEMENT);
+// This statement will fail and cause a rollback as the exception is not caught.
+return transaction.executeUpdate(INVALID_UPDATE_STATEMENT);
+}
+});
+fail("missing expected exception");
+} catch (SpannerException e) {
+assertThat(e.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT);
+}
+assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(0);
+assertThat(countRequests(CommitRequest.class)).isEqualTo(0);
+assertThat(countRequests(ExecuteSqlRequest.class)).isEqualTo(2);
+assertThat(countRequests(RollbackRequest.class)).isEqualTo(1);
+assertThat(countTransactionsStarted()).isEqualTo(1);
+}
+
 @Test
 public void testInlinedBeginTxWithParallelQueries() {
 final int numQueries = 100;

@@ -28,6 +28,7 @@
 import com.google.cloud.spanner.BatchReadOnlyTransaction;
 import com.google.cloud.spanner.Database;
 import com.google.cloud.spanner.DatabaseClient;
+import com.google.cloud.spanner.DatabaseId;
 import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.IntegrationTestEnv;
 import com.google.cloud.spanner.Key;
@@ -45,6 +46,8 @@
 import com.google.cloud.spanner.TimestampBound;
 import com.google.cloud.spanner.TransactionContext;
 import com.google.cloud.spanner.TransactionRunner;
+import com.google.cloud.spanner.TransactionRunner.TransactionCallable;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Sets;
 import com.google.common.util.concurrent.SettableFuture;
 import com.google.common.util.concurrent.Uninterruptibles;
@@ -111,6 +114,7 @@ public void setupClient() {
 
 @After
 public void closeClient() {
+client.writeAtLeastOnce(ImmutableList.of(Mutation.delete("T", KeySet.all())));
 spanner.close();
 }
 
@@ -548,4 +552,68 @@ public Void run(TransactionContext transaction) throws SpannerException {
 }
 });
 }
+
+@Test
+public void testTxWithCaughtError() {
+long updateCount =
+client
+.readWriteTransaction()
+.run(
+new TransactionCallable<Long>() {
+@Override
+public Long run(TransactionContext transaction) throws Exception {
+try {
+transaction.executeUpdate(
+Statement.of("UPDATE NonExistingTable SET Foo=1 WHERE Bar=2"));
+fail("missing expected exception");
+} catch (SpannerException e) {
+assertThat(e.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT);
+}
+return transaction.executeUpdate(
+Statement.of("INSERT INTO T (K, V) VALUES ('One', 1)"));
+}
+});
+assertThat(updateCount).isEqualTo(1L);
+}
+
+@Test
+public void testTxWithUncaughtError() {
+try {
+client
+.readWriteTransaction()
+.run(
+new TransactionCallable<Long>() {
+@Override
+public Long run(TransactionContext transaction) throws Exception {
+return transaction.executeUpdate(
+Statement.of("UPDATE NonExistingTable SET Foo=1 WHERE Bar=2"));
+}
+});
+fail("missing expected exception");
+} catch (SpannerException e) {
+assertThat(e.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT);
+}
+}
+
+@Test
+public void testTxWithUncaughtErrorAfterSuccessfulBegin() {
+DatabaseClient client =
+spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]"));
+try {
+client
+.readWriteTransaction()
+.run(
+new TransactionCallable<Long>() {
+@Override
+public Long run(TransactionContext transaction) throws Exception {
+transaction.executeUpdate(Statement.of("INSERT INTO T (K, V) VALUES ('One', 1)"));
+return transaction.executeUpdate(
+Statement.of("UPDATE NonExistingTable SET Foo=1 WHERE Bar=2"));
+}
+});
+fail("missing expected exception");
+} catch (SpannerException e) {
+assertThat(e.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT);
+}
+}
 }