public Timestamp updateToRunning(String partitionToken) {
    final TransactionResult<Void> transactionResult =
        runInTransaction(
            transaction -> transaction.updateToRunning(partitionToken), "updateToRunning");
    return transactionResult.getCommitTimestamp();
  }
@Test
  public void testInTransactionContextUpdateToRunning() {
    ResultSet resultSet = mock(ResultSet.class);
    when(transaction.executeQuery(any(), anyObject())).thenReturn(resultSet);
    when(resultSet.next()).thenReturn(true);
    when(resultSet.getString(any())).thenReturn(State.SCHEDULED.toString());
    when(resultSet.getCurrentRowAsStruct()).thenReturn(Struct.newBuilder().build());

    ArgumentCaptor<ImmutableList<Mutation>> mutations =
        ArgumentCaptor.forClass(ImmutableList.class);
    doNothing().when(transaction).buffer(mutations.capture());
    assertNull(inTransactionContext.updateToRunning(PARTITION_TOKEN));
    assertEquals(1, mutations.getValue().size());
    Map<String, Value> mutationValueMap = mutations.getValue().iterator().next().asMap();
    assertEquals(
        PARTITION_TOKEN,
        mutationValueMap.get(PartitionMetadataAdminDao.COLUMN_PARTITION_TOKEN).getString());
    assertEquals(
        PartitionMetadata.State.RUNNING.toString(),
        mutationValueMap.get(PartitionMetadataAdminDao.COLUMN_STATE).getString());
  }