public void updateWatermark(String partitionToken, Timestamp watermark) {
    runInTransaction(
        transaction -> transaction.updateWatermark(partitionToken, watermark), "updateWatermark");
  }
@Test
  public void testInTransactionContextUpdateWatermark() {
    ArgumentCaptor<Mutation> mutation = ArgumentCaptor.forClass(Mutation.class);
    doNothing().when(transaction).buffer(mutation.capture());
    assertNull(inTransactionContext.updateWatermark(PARTITION_TOKEN, WATERMARK));
    Map<String, Value> mutationValueMap = mutation.getValue().asMap();
    assertEquals(
        PARTITION_TOKEN,
        mutationValueMap.get(PartitionMetadataAdminDao.COLUMN_PARTITION_TOKEN).getString());
    assertEquals(
        WATERMARK, mutationValueMap.get(PartitionMetadataAdminDao.COLUMN_WATERMARK).getTimestamp());
  }