package storm.kafka; import backtype.storm.spout.SchemeAsMultiScheme; import backtype.storm.utils.Utils; import com.google.common.collect.ImmutableMap; import kafka.api.OffsetRequest; import kafka.javaapi.consumer.SimpleConsumer; import kafka.javaapi.message.ByteBufferMessageSet; import kafka.javaapi.producer.Producer; import kafka.message.MessageAndOffset; import kafka.producer.KeyedMessage; import kafka.producer.ProducerConfig; import org.junit.After; import org.junit.Before; import org.junit.Test; import storm.kafka.trident.GlobalPartitionInformation; import java.util.List; import java.util.Properties; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.is; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertThat; public class KafkaUtilsTest { private KafkaTestBroker broker; private SimpleConsumer simpleConsumer; private KafkaConfig config; private BrokerHosts brokerHosts; @Before public void setup() { broker = new KafkaTestBroker(); GlobalPartitionInformation globalPartitionInformation = new GlobalPartitionInformation(); globalPartitionInformation.addPartition(0, Broker.fromString(broker.getBrokerConnectionString())); brokerHosts = new StaticHosts(globalPartitionInformation); config = new KafkaConfig(brokerHosts, "testTopic"); simpleConsumer = new SimpleConsumer("localhost", broker.getPort(), 60000, 1024, "testClient"); } @After public void shutdown() throws Exception { simpleConsumer.close(); broker.shutdown(); } @Test(expected = FailedFetchException.class) public void topicDoesNotExist() throws Exception { KafkaUtils.fetchMessages(config, simpleConsumer, new Partition(Broker.fromString(broker.getBrokerConnectionString()), 0), 0); } @Test(expected = FailedFetchException.class) public void brokerIsDown() throws Exception { int port = broker.getPort(); broker.shutdown(); SimpleConsumer simpleConsumer = new SimpleConsumer("localhost", port, 100, 1024, "testClient"); try { KafkaUtils.fetchMessages(config, simpleConsumer, new Partition(Broker.fromString(broker.getBrokerConnectionString()), 0), OffsetRequest.LatestTime()); } finally { simpleConsumer.close(); } } @Test public void fetchMessage() throws Exception { String value = "test"; createTopicAndSendMessage(value); long offset = KafkaUtils.getOffset(simpleConsumer, config.topic, 0, OffsetRequest.LatestTime()) - 1; ByteBufferMessageSet messageAndOffsets = KafkaUtils.fetchMessages(config, simpleConsumer, new Partition(Broker.fromString(broker.getBrokerConnectionString()), 0), offset); String message = new String(Utils.toByteArray(messageAndOffsets.iterator().next().message().payload())); assertThat(message, is(equalTo(value))); } @Test(expected = FailedFetchException.class) public void fetchMessagesWithInvalidOffsetAndDefaultHandlingDisabled() throws Exception { config.useStartOffsetTimeIfOffsetOutOfRange = false; KafkaUtils.fetchMessages(config, simpleConsumer, new Partition(Broker.fromString(broker.getBrokerConnectionString()), 0), -99); } @Test public void fetchMessagesWithInvalidOffsetAndDefaultHandlingEnabled() throws Exception { config = new KafkaConfig(brokerHosts, "newTopic"); String value = "test"; createTopicAndSendMessage(value); ByteBufferMessageSet messageAndOffsets = KafkaUtils.fetchMessages(config, simpleConsumer, new Partition(Broker.fromString(broker.getBrokerConnectionString()), 0), -99); String message = new String(Utils.toByteArray(messageAndOffsets.iterator().next().message().payload())); assertThat(message, is(equalTo(value))); } @Test public void getOffsetFromConfigAndDontForceFromStart() { config.forceFromStart = false; config.startOffsetTime = OffsetRequest.EarliestTime(); createTopicAndSendMessage(); long latestOffset = KafkaUtils.getOffset(simpleConsumer, config.topic, 0, OffsetRequest.LatestTime()); long offsetFromConfig = KafkaUtils.getOffset(simpleConsumer, config.topic, 0, config); assertThat(latestOffset, is(equalTo(offsetFromConfig))); } @Test public void getOffsetFromConfigAndFroceFromStart() { config.forceFromStart = true; config.startOffsetTime = OffsetRequest.EarliestTime(); createTopicAndSendMessage(); long earliestOffset = KafkaUtils.getOffset(simpleConsumer, config.topic, 0, OffsetRequest.EarliestTime()); long offsetFromConfig = KafkaUtils.getOffset(simpleConsumer, config.topic, 0, config); assertThat(earliestOffset, is(equalTo(offsetFromConfig))); } @Test public void generateTuplesWithoutKeyAndKeyValueScheme() { config.scheme = new KeyValueSchemeAsMultiScheme(new StringKeyValueScheme()); runGetValueOnlyTuplesTest(); } @Test public void generateTuplesWithKeyAndKeyValueScheme() { config.scheme = new KeyValueSchemeAsMultiScheme(new StringKeyValueScheme()); String value = "value"; String key = "key"; createTopicAndSendMessage(key, value); ByteBufferMessageSet messageAndOffsets = getLastMessage(); for (MessageAndOffset msg : messageAndOffsets) { Iterable<List<Object>> lists = KafkaUtils.generateTuples(config, msg.message()); assertEquals(ImmutableMap.of(key, value), lists.iterator().next().get(0)); } } @Test public void generateTupelsWithValueScheme() { config.scheme = new SchemeAsMultiScheme(new StringScheme()); runGetValueOnlyTuplesTest(); } @Test public void generateTuplesWithValueSchemeAndKeyValueMessage() { config.scheme = new SchemeAsMultiScheme(new StringScheme()); String value = "value"; String key = "key"; createTopicAndSendMessage(key, value); ByteBufferMessageSet messageAndOffsets = getLastMessage(); for (MessageAndOffset msg : messageAndOffsets) { Iterable<List<Object>> lists = KafkaUtils.generateTuples(config, msg.message()); assertEquals(value, lists.iterator().next().get(0)); } } private ByteBufferMessageSet getLastMessage() { long offsetOfLastMessage = KafkaUtils.getOffset(simpleConsumer, config.topic, 0, OffsetRequest.LatestTime()) - 1; return KafkaUtils.fetchMessages(config, simpleConsumer, new Partition(Broker.fromString(broker.getBrokerConnectionString()), 0), offsetOfLastMessage); } private void runGetValueOnlyTuplesTest() { String value = "value"; createTopicAndSendMessage(null, value); ByteBufferMessageSet messageAndOffsets = getLastMessage(); for (MessageAndOffset msg : messageAndOffsets) { Iterable<List<Object>> lists = KafkaUtils.generateTuples(config, msg.message()); assertEquals(value, lists.iterator().next().get(0)); } } private void createTopicAndSendMessage() { createTopicAndSendMessage(null, "someValue"); } private void createTopicAndSendMessage(String value) { createTopicAndSendMessage(null, value); } private void createTopicAndSendMessage(String key, String value) { Properties p = new Properties(); p.setProperty("metadata.broker.list", broker.getBrokerConnectionString()); p.setProperty("serializer.class", "kafka.serializer.StringEncoder"); ProducerConfig producerConfig = new ProducerConfig(p); Producer<String, String> producer = new Producer<String, String>(producerConfig); producer.send(new KeyedMessage<String, String>(config.topic, key, value)); } @Test public void assignOnePartitionPerTask() { runPartitionToTaskMappingTest(16, 1); } @Test public void assignTwoPartitionsPerTask() { runPartitionToTaskMappingTest(16, 2); } @Test public void assignAllPartitionsToOneTask() { runPartitionToTaskMappingTest(32, 32); } public void runPartitionToTaskMappingTest(int numPartitions, int partitionsPerTask) { GlobalPartitionInformation globalPartitionInformation = TestUtils.buildPartitionInfo(numPartitions); int numTasks = numPartitions / partitionsPerTask; for (int i = 0 ; i < numTasks ; i++) { assertEquals(partitionsPerTask, KafkaUtils.calculatePartitionsForTask(globalPartitionInformation, numTasks, i).size()); } } @Test public void moreTasksThanPartitions() { GlobalPartitionInformation globalPartitionInformation = TestUtils.buildPartitionInfo(1); int numTasks = 2; assertEquals(1, KafkaUtils.calculatePartitionsForTask(globalPartitionInformation, numTasks, 0).size()); assertEquals(0, KafkaUtils.calculatePartitionsForTask(globalPartitionInformation, numTasks, 1).size()); } @Test (expected = IllegalArgumentException.class ) public void assignInvalidTask() { GlobalPartitionInformation globalPartitionInformation = new GlobalPartitionInformation(); KafkaUtils.calculatePartitionsForTask(globalPartitionInformation, 1, 1); } }