use of com.amazonaws.services.kinesis.clientlibrary.interfaces.v2.IRecordProcessorFactory in project samza by apache.
the class TestKinesisSystemConsumer method testProcessRecordsHelper.
/**
* Helper to simulate and test the life-cycle of record processing from a kinesis stream with a given number of shards
* 1. Creation of record processors.
* 2. Initialization of record processors.
* 3. Processing records via record processors.
* 4. Calling checkpoint on record processors.
* 5. Shutting down (due to re-assignment or lease expiration) record processors.
*/
private void testProcessRecordsHelper(String system, String stream, int numShards, int numRecordsPerShard) throws InterruptedException, NoSuchFieldException, IllegalAccessException {
KinesisConfig kConfig = new KinesisConfig(new MapConfig());
// Create consumer
KinesisSystemConsumer consumer = new KinesisSystemConsumer(system, kConfig, new NoOpMetricsRegistry());
initializeMetrics(consumer, stream);
List<SystemStreamPartition> ssps = new LinkedList<>();
IntStream.range(0, numShards).forEach(p -> {
SystemStreamPartition ssp = new SystemStreamPartition(system, stream, new Partition(p));
ssps.add(ssp);
});
ssps.forEach(ssp -> consumer.register(ssp, SYSTEM_CONSUMER_REGISTER_OFFSET));
// Create Kinesis record processor factory
IRecordProcessorFactory factory = consumer.createRecordProcessorFactory(stream);
// Create and initialize Kinesis record processor
Map<String, KinesisRecordProcessor> processorMap = createAndInitProcessors(factory, numShards);
List<KinesisRecordProcessor> processorList = new ArrayList<>(processorMap.values());
// Generate records to Kinesis record processor
Map<KinesisRecordProcessor, List<Record>> inputRecordMap = generateRecords(numRecordsPerShard, processorList);
// Verification steps
// Read events from the BEM queue
Map<SystemStreamPartition, List<IncomingMessageEnvelope>> messages = readEvents(new HashSet<>(ssps), consumer, numRecordsPerShard);
if (numRecordsPerShard > 0) {
Assert.assertEquals(messages.size(), numShards);
} else {
// No input records and hence no messages
Assert.assertEquals(messages.size(), 0);
return;
}
Map<SystemStreamPartition, KinesisRecordProcessor> sspToProcessorMap = getProcessorMap(consumer);
ssps.forEach(ssp -> {
try {
KinesisRecordProcessor processor = sspToProcessorMap.get(ssp);
// Verify that the read messages are received in order and are the same as input records
Assert.assertEquals(messages.get(ssp).size(), numRecordsPerShard);
List<IncomingMessageEnvelope> envelopes = messages.get(ssp);
List<Record> inputRecords = inputRecordMap.get(processor);
verifyRecords(envelopes, inputRecords, processor.getShardId());
// Call checkpoint on consumer and verify that the checkpoint is called with the right offset
IncomingMessageEnvelope lastEnvelope = envelopes.get(envelopes.size() - 1);
consumer.afterCheckpoint(Collections.singletonMap(ssp, lastEnvelope.getOffset()));
ArgumentCaptor<String> argument = ArgumentCaptor.forClass(String.class);
verify(getCheckpointer(processor)).checkpoint(argument.capture());
Assert.assertEquals(inputRecords.get(inputRecords.size() - 1).getSequenceNumber(), argument.getValue());
// Call shutdown (with ZOMBIE reason) on processor and verify if shutdown freed the ssp mapping
shutDownProcessor(processor, ShutdownReason.ZOMBIE);
Assert.assertFalse(sspToProcessorMap.containsValue(processor));
Assert.assertTrue(isSspAvailable(consumer, ssp));
} catch (NoSuchFieldException | IllegalAccessException | InvalidStateException | ShutdownException ex) {
throw new RuntimeException(ex);
}
});
}
Aggregations