Examples with TransactionContextImpl com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl used on opensource projects

Example 11 with TransactionContextImpl

use of com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl in project java-spanner by googleapis.

the class TransactionRunnerImplTest method batchDmlException.

@SuppressWarnings("unchecked")
private long[] batchDmlException(int status) {
    Preconditions.checkArgument(status != Code.OK_VALUE);
    TransactionContextImpl transaction = TransactionContextImpl.newBuilder().setSession(session).setTransactionId(ByteString.copyFromUtf8(UUID.randomUUID().toString())).setOptions(Options.fromTransactionOptions()).setRpc(rpc).build();
    when(session.newTransaction(Options.fromTransactionOptions())).thenReturn(transaction);
    when(session.beginTransactionAsync()).thenReturn(ApiFutures.immediateFuture(ByteString.copyFromUtf8(UUID.randomUUID().toString())));
    when(session.getName()).thenReturn(SessionId.of("p", "i", "d", "test").getName());
    TransactionRunnerImpl runner = new TransactionRunnerImpl(session);
    runner.setSpan(mock(Span.class));
    ExecuteBatchDmlResponse response1 = ExecuteBatchDmlResponse.newBuilder().addResultSets(ResultSet.newBuilder().setStats(ResultSetStats.newBuilder().setRowCountExact(1L)).build()).setStatus(com.google.rpc.Status.newBuilder().setCode(status).build()).build();
    ExecuteBatchDmlResponse response2 = ExecuteBatchDmlResponse.newBuilder().addResultSets(ResultSet.newBuilder().setStats(ResultSetStats.newBuilder().setRowCountExact(1L)).build()).addResultSets(ResultSet.newBuilder().setStats(ResultSetStats.newBuilder().setRowCountExact(1L)).build()).setStatus(com.google.rpc.Status.newBuilder().setCode(Code.OK_VALUE).build()).build();
    when(rpc.executeBatchDml(Mockito.any(ExecuteBatchDmlRequest.class), Mockito.anyMap())).thenReturn(response1, response2);
    CommitResponse commitResponse = CommitResponse.newBuilder().setCommitTimestamp(Timestamp.getDefaultInstance()).build();
    when(rpc.commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap())).thenReturn(ApiFutures.immediateFuture(commitResponse));
    final Statement statement = Statement.of("UPDATE FOO SET BAR=1");
    final AtomicInteger numCalls = new AtomicInteger(0);
    long[] updateCount = runner.run(transaction1 -> {
        numCalls.incrementAndGet();
        return transaction1.batchUpdate(Arrays.asList(statement, statement));
    });
    if (status == Code.ABORTED_VALUE) {
        // Assert that the method ran twice because the first response aborted.
        assertThat(numCalls.get()).isEqualTo(2);
    }
    return updateCount;
}