@Override public void apply() throws Exception { // consuming and aborting the partition numRetries times plus one (for the first attempt) makes it get removed // from the working set for (int i = 0; i < numRetries + 1; i++) { List<PartitionDetail> partitionDetails = partitionConsumer.consumePartitions(1).getPartitions(); Assert.assertEquals(1, partitionDetails.size()); Assert.assertEquals(partitionKey, partitionDetails.get(0).getPartitionKey()); // aborting the processing of the partition partitionConsumer.onFinish(partitionDetails, false); } // after the 2nd abort, the partition is discarded entirely, and so no partitions are available for consuming PartitionConsumerResult result = partitionConsumer.consumePartitions(1); Assert.assertEquals(0, result.getPartitions().size()); Assert.assertEquals(1, result.getFailedPartitions().size()); Assert.assertEquals(partitionKey, result.getFailedPartitions().get(0).getPartitionKey()); } });
@Override public void apply() throws Exception { // first call to consume will drop the partition from the working set, and return nothing, since it was // the only partition in the working set PartitionConsumerResult result = partitionConsumer.consumePartitions(1); Assert.assertEquals(0, result.getPartitions().size()); Assert.assertEquals(0, result.getFailedPartitions().size()); // following calls to consumePartitions will repopulate the working set and return additional partition(s) result = partitionConsumer.consumePartitions(1); Assert.assertEquals(1, result.getPartitions().size()); Assert.assertEquals(partitionKey2, result.getPartitions().get(0).getPartitionKey()); } });
@Override public void apply() throws Exception { // consume all the partitions List<? extends Partition> consumedPartitions = partitionConsumer.consumePartitions().getPartitions(); Assert.assertEquals(partitionKeys, toKeys(consumedPartitions)); // consuming the partitions again, without adding any new partitions returns an empty iterator Assert.assertTrue(partitionConsumer.consumePartitions().getPartitions().isEmpty()); // we configured the max number of retries to be 1. However, we are putting back all the partitions 5 times, // and testing that they are still available for processing, and that there are no failed partitions for (int i = 0; i < 5; i++) { partitionConsumer.untake(consumedPartitions); PartitionConsumerResult result = partitionConsumer.consumePartitions(); consumedPartitions = result.getPartitions(); Assert.assertEquals(partitionKeys, toKeys(consumedPartitions)); Assert.assertEquals(0, result.getFailedPartitions().size()); } // consuming the partitions again, without adding any new partitions returns an empty iterator Assert.assertTrue(partitionConsumer.consumePartitions().getPartitions().isEmpty()); // test functionality to put back a partial subset of the retrieved the partitions Partition firstConsumedPartition = consumedPartitions.get(0); // test the untakeWithKeys method partitionConsumer.untakeWithKeys(ImmutableList.of(firstConsumedPartition.getPartitionKey())); consumedPartitions = partitionConsumer.consumePartitions().getPartitions(); Assert.assertEquals(1, consumedPartitions.size()); Assert.assertEquals(firstConsumedPartition, consumedPartitions.get(0)); } });