refactor tests

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

@@ -226,12 +226,10 @@ public ServerStream<BatchWriteResponse> batchWriteAtleastOnceWithOptions(
     }
     Span span = tracer.spanBuilder(SpannerImpl.BATCH_WRITE).startSpan();
     try (Scope s = tracer.withSpan(span)) {
-      ServerStream<BatchWriteResponse> response =
-          spanner.getRpc().batchWriteAtleastOnce(requestBuilder.build(), this.options);
-      return response;
-    } catch (RuntimeException e) {
+      return spanner.getRpc().batchWriteAtleastOnce(requestBuilder.build(), this.options);
+    } catch (Throwable e) {
       TraceUtil.setWithFailure(span, e);
-      throw e;
+      throw SpannerExceptionFactory.newSpannerException(e);
     } finally {
       span.end(TraceUtil.END_SPAN_OPTIONS);
     }

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

@@ -29,6 +29,7 @@
 import static org.junit.Assume.assumeFalse;
 import static org.junit.Assume.assumeTrue;
 
+import com.google.api.gax.rpc.ServerStream;
 import com.google.cloud.ByteArray;
 import com.google.cloud.Date;
 import com.google.cloud.Timestamp;
@@ -53,6 +54,9 @@
 import com.google.cloud.spanner.testing.EmulatorSpannerHelper;
 import com.google.common.collect.ImmutableList;
 import com.google.protobuf.NullValue;
+import com.google.rpc.Code;
+import com.google.rpc.Status;
+import com.google.spanner.v1.BatchWriteResponse;
 import io.grpc.Context;
 import java.math.BigDecimal;
 import java.util.ArrayList;
@@ -212,6 +216,28 @@ public void writeAtLeastOnce() {
     assertThat(row.getString(0)).isEqualTo("v1");
   }
 
+  @Test
+  public void batchWriteAtLeastOnce() {
+    ServerStream<BatchWriteResponse> responses =
+        client.batchWriteAtleastOnce(
+            ImmutableList.of(
+                Mutation.newInsertOrUpdateBuilder("T")
+                    .set("K")
+                    .to(lastKey = uniqueString())
+                    .set("StringValue")
+                    .to("v1")
+                    .build()));
+    for (BatchWriteResponse response : responses) {
+      assertEquals(response.getIndexesList(), Collections.singletonList(0));
+      if (response.getStatus().equals(Status.newBuilder().setCode(Code.OK_VALUE).build())) {
+        assertNotNull(response.getCommitTimestamp());
+        Struct row = readLastRow("StringValue");
+        assertFalse(row.isNull(0));
+        assertEquals(row.getString(0), "v1");
+      }
+    }
+  }
+
   @Test
   public void testWriteReturnsCommitStats() {
     assumeFalse("Emulator does not return commit statistics", isUsingEmulator());