From c20ef851687f396d4b2e9f48db351f3ea2b9b97b Mon Sep 17 00:00:00 2001 From: Anshul Mehra Date: Wed, 28 Aug 2019 15:55:42 -0400 Subject: [PATCH] GH-280: Add topicPattern support to KafkaMessageSource Fixes https://github.com/spring-projects/spring-integration-kafka/issues/280 * Refactored to topics/topicPattern to be a property to prevent constructor telescoping * Add support for manual partition assignment - Cleanup after rebasing from #283 - Add additional tests * Fix checkstyle * Remove redundant factory methods * Add better test for static assignment * Use ConsumerProperties Fix typo where ContainerProperties was used instead of ConsumerProperties * Simplify construction * Remove unused import * Use ObjectUtils * Use component name for clientId * Collapse static assignment for loop * Fix lint warning --- .../kafka/inbound/KafkaMessageSource.java | 152 +++++++--- ...KafkaInboundChannelAdapterParserTests.java | 10 +- .../kafka/inbound/MessageSourceTests.java | 283 ++++++++++++++++++ 3 files changed, 392 insertions(+), 53 deletions(-) diff --git a/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/inbound/KafkaMessageSource.java b/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/inbound/KafkaMessageSource.java index 361c1f027b..3af243b835 100644 --- a/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/inbound/KafkaMessageSource.java +++ b/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/inbound/KafkaMessageSource.java @@ -55,12 +55,14 @@ import org.springframework.kafka.listener.ConsumerAwareRebalanceListener; import org.springframework.kafka.listener.ConsumerProperties; import org.springframework.kafka.support.Acknowledgment; import org.springframework.kafka.support.KafkaHeaders; +import org.springframework.kafka.support.TopicPartitionOffset; import org.springframework.kafka.support.converter.KafkaMessageHeaders; import org.springframework.kafka.support.converter.MessagingMessageConverter; import org.springframework.kafka.support.converter.RecordMessageConverter; import org.springframework.lang.Nullable; import org.springframework.messaging.Message; import org.springframework.util.Assert; +import org.springframework.util.ObjectUtils; import org.springframework.util.StringUtils; /** @@ -107,17 +109,13 @@ public class KafkaMessageSource extends AbstractMessageSource impl private final KafkaAckCallbackFactory ackCallbackFactory; - private final String[] topics; - private final Object consumerMonitor = new Object(); private final Map>> inflightRecords = new ConcurrentHashMap<>(); private final AtomicInteger remainingCount = new AtomicInteger(); - private String groupId; - - private String clientId = "message.source"; + private final ConsumerProperties consumerProperties; private Duration pollTimeout; @@ -125,10 +123,6 @@ public class KafkaMessageSource extends AbstractMessageSource impl private Class payloadType; - private ConsumerRebalanceListener rebalanceListener; - - private ConsumerAwareRebalanceListener consumerAwareRebalanceListener; - private boolean rawMessageHeader; private Duration commitTimeout; @@ -250,28 +244,31 @@ public class KafkaMessageSource extends AbstractMessageSource impl Assert.notNull(consumerFactory, "'consumerFactory' must not be null"); Assert.notNull(ackCallbackFactory, "'ackCallbackFactory' must not be null"); - Assert.isTrue(consumerProperties.getTopics() != null && consumerProperties.getTopics().length > 0, "At least one topic is required"); + Assert.isTrue( + !ObjectUtils.isEmpty(consumerProperties.getTopics()) + || !ObjectUtils.isEmpty(consumerProperties.getTopicPartitionsToAssign()) + || consumerProperties.getTopicPattern() != null, + "topics, topicPattern, or topicPartitions must be provided" + ); + this.consumerProperties = consumerProperties; + this.consumerFactory = fixOrRejectConsumerFactory(consumerFactory, allowMultiFetch); this.ackCallbackFactory = ackCallbackFactory; - this.topics = consumerProperties.getTopics(); - this.groupId = consumerProperties.getGroupId(); - if (StringUtils.hasText(consumerProperties.getClientId())) { - this.clientId = consumerProperties.getClientId(); - } this.pollTimeout = Duration.ofMillis(consumerProperties.getPollTimeout()); this.assignTimeout = this.minTimeoutProvider.get(); this.commitTimeout = consumerProperties.getSyncCommitTimeout(); this.ackCallbackFactory.setCommitTimeout(consumerProperties.getSyncCommitTimeout()); - if (consumerProperties.getConsumerRebalanceListener() instanceof ConsumerAwareRebalanceListener) { - this.consumerAwareRebalanceListener = (ConsumerAwareRebalanceListener) consumerProperties.getConsumerRebalanceListener(); - } - else { - this.rebalanceListener = consumerProperties.getConsumerRebalanceListener(); + } + + @Override + protected void onInit() { + if (!StringUtils.hasText(this.consumerProperties.getClientId())) { + this.consumerProperties.setClientId(getComponentName()); } } protected String getGroupId() { - return this.groupId; + return this.consumerProperties.getGroupId(); } /** @@ -282,11 +279,11 @@ public class KafkaMessageSource extends AbstractMessageSource impl */ @Deprecated public void setGroupId(String groupId) { - this.groupId = groupId; + this.consumerProperties.setGroupId(groupId); } protected String getClientId() { - return this.clientId; + return this.consumerProperties.getClientId(); } /** @@ -297,7 +294,7 @@ public class KafkaMessageSource extends AbstractMessageSource impl */ @Deprecated public void setClientId(String clientId) { - this.clientId = clientId; + this.consumerProperties.setClientId(clientId); } protected long getPollTimeout() { @@ -343,7 +340,7 @@ public class KafkaMessageSource extends AbstractMessageSource impl } protected ConsumerRebalanceListener getRebalanceListener() { - return this.rebalanceListener; + return this.consumerProperties.getConsumerRebalanceListener(); } /** @@ -354,10 +351,7 @@ public class KafkaMessageSource extends AbstractMessageSource impl */ @Deprecated public void setRebalanceListener(ConsumerRebalanceListener rebalanceListener) { - this.rebalanceListener = rebalanceListener; - if (rebalanceListener instanceof ConsumerAwareRebalanceListener) { - this.consumerAwareRebalanceListener = (ConsumerAwareRebalanceListener) rebalanceListener; - } + this.consumerProperties.setConsumerRebalanceListener(rebalanceListener); } @Override @@ -537,41 +531,103 @@ public class KafkaMessageSource extends AbstractMessageSource impl protected void createConsumer() { synchronized (this.consumerMonitor) { - this.consumer = this.consumerFactory.createConsumer(this.groupId, this.clientId, null); - boolean isConsumerAware = this.consumerAwareRebalanceListener != null; - this.consumer.subscribe(Arrays.asList(this.topics), new ConsumerRebalanceListener() { + this.consumer = this.consumerFactory.createConsumer(this.consumerProperties.getGroupId(), + this.consumerProperties.getClientId(), null); + ConsumerRebalanceListener providedRebalanceListener = this.consumerProperties + .getConsumerRebalanceListener(); + boolean isConsumerAware = providedRebalanceListener instanceof ConsumerAwareRebalanceListener; + ConsumerRebalanceListener rebalanceCallback = new ConsumerRebalanceListener() { @Override public void onPartitionsRevoked(Collection partitions) { KafkaMessageSource.this.assignedPartitions.clear(); if (KafkaMessageSource.this.logger.isInfoEnabled()) { - KafkaMessageSource.this.logger.info("Partitions revoked: " + partitions); + KafkaMessageSource.this.logger + .info("Partitions revoked: " + partitions); } - if (isConsumerAware) { - KafkaMessageSource.this.consumerAwareRebalanceListener.onPartitionsRevokedAfterCommit( - KafkaMessageSource.this.consumer, partitions); - } - else if (KafkaMessageSource.this.rebalanceListener != null) { - KafkaMessageSource.this.rebalanceListener.onPartitionsRevoked(partitions); + if (providedRebalanceListener != null) { + if (isConsumerAware) { + ((ConsumerAwareRebalanceListener) providedRebalanceListener) + .onPartitionsRevokedAfterCommit(KafkaMessageSource.this.consumer, partitions); + } + else { + providedRebalanceListener.onPartitionsRevoked(partitions); + } } + } @Override public void onPartitionsAssigned(Collection partitions) { KafkaMessageSource.this.assignedPartitions = new ArrayList<>(partitions); if (KafkaMessageSource.this.logger.isInfoEnabled()) { - KafkaMessageSource.this.logger.info("Partitions assigned: " + partitions); + KafkaMessageSource.this.logger + .info("Partitions assigned: " + partitions); } - if (isConsumerAware) { - KafkaMessageSource.this.consumerAwareRebalanceListener.onPartitionsAssigned( - KafkaMessageSource.this.consumer, partitions); - } - else if (KafkaMessageSource.this.rebalanceListener != null) { - KafkaMessageSource.this.rebalanceListener.onPartitionsAssigned(partitions); + if (providedRebalanceListener != null) { + if (isConsumerAware) { + ((ConsumerAwareRebalanceListener) providedRebalanceListener) + .onPartitionsAssigned(KafkaMessageSource.this.consumer, partitions); + } + else { + providedRebalanceListener.onPartitionsAssigned(partitions); + } } } - }); + }; + + if (this.consumerProperties.getTopicPattern() != null) { + this.consumer.subscribe(this.consumerProperties.getTopicPattern(), rebalanceCallback); + } + else if (this.consumerProperties.getTopicPartitionsToAssign() != null) { + List topicPartitionsToAssign = Arrays + .stream(this.consumerProperties.getTopicPartitionsToAssign()) + .map(TopicPartitionOffset::getTopicPartition) + .collect(Collectors.toList()); + this.consumer.assign(topicPartitionsToAssign); + this.assignedPartitions = new ArrayList<>(topicPartitionsToAssign); + + TopicPartitionOffset[] partitions = this.consumerProperties.getTopicPartitionsToAssign(); + + for (TopicPartitionOffset partition : partitions) { + if (TopicPartitionOffset.SeekPosition.BEGINNING.equals(partition.getPosition())) { + this.consumer.seekToBeginning(Collections.singleton(partition.getTopicPartition())); + } + else if (TopicPartitionOffset.SeekPosition.END.equals(partition.getPosition())) { + this.consumer.seekToEnd(Collections.singleton(partition.getTopicPartition())); + } + else { + TopicPartition topicPartition = partition.getTopicPartition(); + Long offset = partition.getOffset(); + if (offset != null) { + long newOffset = offset; + + if (offset < 0) { + if (!partition.isRelativeToCurrent()) { + this.consumer.seekToEnd(Collections.singleton(topicPartition)); + continue; + } + newOffset = Math.max(0, this.consumer.position(topicPartition) + offset); + } + else if (partition.isRelativeToCurrent()) { + newOffset = this.consumer.position(topicPartition) + offset; + } + + try { + this.consumer.seek(topicPartition, newOffset); + } + catch (Exception e) { + this.logger.error("Failed to set initial offset for " + topicPartition + + " at " + newOffset + ". Position is " + this.consumer.position(topicPartition), e); + } + } + } + } + } + else { + this.consumer.subscribe(Arrays.asList(this.consumerProperties.getTopics()), rebalanceCallback); + } } } @@ -799,7 +855,7 @@ public class KafkaMessageSource extends AbstractMessageSource impl @Override public String getGroupId() { - return KafkaMessageSource.this.groupId; + return KafkaMessageSource.this.getGroupId(); } @Override diff --git a/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/config/xml/KafkaInboundChannelAdapterParserTests.java b/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/config/xml/KafkaInboundChannelAdapterParserTests.java index d5ca499d9c..7bcbf14554 100644 --- a/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/config/xml/KafkaInboundChannelAdapterParserTests.java +++ b/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/config/xml/KafkaInboundChannelAdapterParserTests.java @@ -52,21 +52,21 @@ public class KafkaInboundChannelAdapterParserTests { @Test public void testProps() { - assertThat(TestUtils.getPropertyValue(this.source1, "topics")).isEqualTo(new String[] { "topic1" }); + assertThat(TestUtils.getPropertyValue(this.source1, "consumerProperties.topics")).isEqualTo(new String[] { "topic1" }); assertThat(TestUtils.getPropertyValue(this.source1, "consumerFactory")) .isSameAs(this.context.getBean("consumerFactory")); assertThat(TestUtils.getPropertyValue(this.source1, "ackCallbackFactory")) .isSameAs(this.context.getBean("ackFactory")); - assertThat(TestUtils.getPropertyValue(this.source1, "clientId")).isEqualTo("client"); - assertThat(TestUtils.getPropertyValue(this.source1, "groupId")).isEqualTo("group"); + assertThat(TestUtils.getPropertyValue(this.source1, "consumerProperties.clientId")).isEqualTo("client"); + assertThat(TestUtils.getPropertyValue(this.source1, "consumerProperties.groupId")).isEqualTo("group"); assertThat(TestUtils.getPropertyValue(this.source1, "messageConverter")) .isSameAs(this.context.getBean("converter")); assertThat(TestUtils.getPropertyValue(this.source1, "payloadType")).isEqualTo(String.class); assertThat(TestUtils.getPropertyValue(this.source1, "rawMessageHeader", Boolean.class)).isTrue(); - assertThat(TestUtils.getPropertyValue(this.source1, "rebalanceListener")) + assertThat(TestUtils.getPropertyValue(this.source1, "consumerProperties.consumerRebalanceListener")) .isSameAs(this.context.getBean("rebal")); - assertThat(TestUtils.getPropertyValue(this.source2, "topics")).isEqualTo(new String[] { "topic1", "topic2" }); + assertThat(TestUtils.getPropertyValue(this.source2, "consumerProperties.topics")).isEqualTo(new String[] { "topic1", "topic2" }); DefaultKafkaConsumerFactory cf = TestUtils.getPropertyValue(this.source2, "consumerFactory", DefaultKafkaConsumerFactory.class); assertThat(cf).isSameAs(this.context.getBean("multiFetchConsumerFactory")); diff --git a/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageSourceTests.java b/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageSourceTests.java index 7793dda65c..2d3a8c2f36 100644 --- a/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageSourceTests.java +++ b/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageSourceTests.java @@ -17,9 +17,11 @@ package org.springframework.integration.kafka.inbound; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.assertj.core.api.Assertions.fail; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyCollection; +import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.isNull; import static org.mockito.BDDMockito.given; @@ -29,6 +31,7 @@ import static org.mockito.BDDMockito.willReturn; import static org.mockito.Mockito.inOrder; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.times; import java.time.Duration; import java.time.temporal.ChronoUnit; @@ -40,7 +43,13 @@ import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Set; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; +import java.util.regex.Pattern; +import java.util.stream.Collectors; +import java.util.stream.Stream; import org.apache.commons.logging.Log; import org.apache.kafka.clients.consumer.Consumer; @@ -48,7 +57,9 @@ import org.apache.kafka.clients.consumer.ConsumerConfig; import org.apache.kafka.clients.consumer.ConsumerRebalanceListener; import org.apache.kafka.clients.consumer.ConsumerRecord; import org.apache.kafka.clients.consumer.ConsumerRecords; +import org.apache.kafka.clients.consumer.MockConsumer; import org.apache.kafka.clients.consumer.OffsetAndMetadata; +import org.apache.kafka.clients.consumer.OffsetResetStrategy; import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.record.TimestampType; import org.junit.Test; @@ -62,8 +73,10 @@ import org.springframework.integration.acks.AcknowledgmentCallback; import org.springframework.integration.test.util.TestUtils; import org.springframework.kafka.core.ConsumerFactory; import org.springframework.kafka.core.DefaultKafkaConsumerFactory; +import org.springframework.kafka.listener.ConsumerAwareRebalanceListener; import org.springframework.kafka.listener.ConsumerProperties; import org.springframework.kafka.support.KafkaHeaders; +import org.springframework.kafka.support.TopicPartitionOffset; import org.springframework.kafka.test.utils.KafkaTestUtils; import org.springframework.messaging.Message; @@ -76,6 +89,102 @@ import org.springframework.messaging.Message; */ public class MessageSourceTests { + @Test + public void testIllegalArgs() { + ConsumerFactory consumerFactory = mock(ConsumerFactory.class); + assertThatThrownBy(() -> new KafkaMessageSource(consumerFactory, new ConsumerProperties((Pattern) null))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("topics, topicPattern, or topicPartitions must be provided"); + } + + @Test + public void testConsumerAwareRebalanceListener() { + Consumer consumer = mock(Consumer.class); + TopicPartition topicPartition = new TopicPartition("foo", 0); + List assigned = Collections.singletonList(topicPartition); + AtomicReference listener = new AtomicReference<>(); + willAnswer(i -> { + listener.set(i.getArgument(1)); + return null; + }).given(consumer).subscribe(anyCollection(), any(ConsumerRebalanceListener.class)); + + ConsumerFactory consumerFactory = mock(ConsumerFactory.class); + willReturn(Collections.singletonMap(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, 1)).given(consumerFactory) + .getConfigurationProperties(); + given(consumerFactory.createConsumer(isNull(), anyString(), isNull())).willReturn(consumer); + ConsumerProperties consumerProperties = new ConsumerProperties("foo"); + AtomicBoolean partitionsAssignedCalled = new AtomicBoolean(); + AtomicReference partitionsAssignedConsumer = new AtomicReference<>(); + AtomicBoolean partitionsRevokedCalled = new AtomicBoolean(); + AtomicReference partitionsRevokedConsumer = new AtomicReference<>(); + consumerProperties.setConsumerRebalanceListener(new ConsumerAwareRebalanceListener() { + @Override + public void onPartitionsRevokedAfterCommit(Consumer consumer, Collection partitions) { + partitionsRevokedCalled.getAndSet(true); + partitionsRevokedConsumer.set(consumer); + } + + @Override + public void onPartitionsAssigned(Consumer consumer, Collection partitions) { + partitionsAssignedCalled.getAndSet(true); + partitionsAssignedConsumer.set(consumer); + } + }); + KafkaMessageSource source = new KafkaMessageSource(consumerFactory, consumerProperties); + source.setRawMessageHeader(true); + + Message received = source.receive(); + + listener.get().onPartitionsAssigned(assigned); + assertThat(partitionsAssignedCalled.get()).isTrue(); + assertThat(partitionsAssignedConsumer.get()).isEqualTo(consumer); + + listener.get().onPartitionsRevoked(assigned); + assertThat(partitionsRevokedCalled.get()).isTrue(); + assertThat(partitionsRevokedConsumer.get()).isEqualTo(consumer); + } + + @Test + public void testRebalanceListener() { + Consumer consumer = mock(Consumer.class); + TopicPartition topicPartition = new TopicPartition("foo", 0); + List assigned = Collections.singletonList(topicPartition); + AtomicReference listener = new AtomicReference<>(); + willAnswer(i -> { + listener.set(i.getArgument(1)); + return null; + }).given(consumer).subscribe(anyCollection(), any(ConsumerRebalanceListener.class)); + + ConsumerFactory consumerFactory = mock(ConsumerFactory.class); + willReturn(Collections.singletonMap(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, 1)).given(consumerFactory) + .getConfigurationProperties(); + given(consumerFactory.createConsumer(isNull(), anyString(), isNull())).willReturn(consumer); + ConsumerProperties consumerProperties = new ConsumerProperties("foo"); + AtomicBoolean partitionsAssignedCalled = new AtomicBoolean(); + AtomicBoolean partitionsRevokedCalled = new AtomicBoolean(); + consumerProperties.setConsumerRebalanceListener(new ConsumerRebalanceListener() { + @Override + public void onPartitionsRevoked(Collection partitions) { + partitionsRevokedCalled.getAndSet(true); + } + + @Override + public void onPartitionsAssigned(Collection partitions) { + partitionsAssignedCalled.getAndSet(true); + } + }); + KafkaMessageSource source = new KafkaMessageSource(consumerFactory, consumerProperties); + source.setRawMessageHeader(true); + + Message received = source.receive(); + + listener.get().onPartitionsAssigned(assigned); + assertThat(partitionsAssignedCalled.get()).isTrue(); + + listener.get().onPartitionsRevoked(assigned); + assertThat(partitionsRevokedCalled.get()).isTrue(); + } + @SuppressWarnings({ "rawtypes", "unchecked" }) @Test public void testAck() { @@ -553,4 +662,178 @@ public class MessageSourceTests { inOrder.verifyNoMoreInteractions(); } + @SuppressWarnings("unchecked") + @Test + public void testTopicPatternBasedMessageSource() { + MockConsumer consumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST); + TopicPartition topicPartition1 = new TopicPartition("abc_foo", 0); + TopicPartition topicPartition2 = new TopicPartition("abc_foo", 1); + TopicPartition topicPartition3 = new TopicPartition("def_foo", 0); + TopicPartition topicPartition4 = new TopicPartition("def_foo", 1); + List topicPartitions = Arrays + .asList(topicPartition1, topicPartition2, topicPartition3, topicPartition4); + + Map beginningOffsets = topicPartitions.stream().collect(Collectors + .toMap(Function.identity(), tp -> 0L)); + consumer.updateBeginningOffsets(beginningOffsets); + + ConsumerFactory consumerFactory = mock(ConsumerFactory.class); + willReturn(Collections.singletonMap(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, 1)).given(consumerFactory) + .getConfigurationProperties(); + given(consumerFactory.createConsumer(isNull(), anyString(), isNull())).willReturn(consumer); + KafkaMessageSource source = new KafkaMessageSource<>(consumerFactory, new ConsumerProperties(Pattern.compile("[a-zA-Z0-9_]*?foo"))); + source.setRawMessageHeader(true); + source.start(); + // force consumer creation + source.receive(); + + consumer.rebalance(topicPartitions); + + ConsumerRecord record1 = new ConsumerRecord<>("abc_foo", 0, 0, null, "a"); + ConsumerRecord record2 = new ConsumerRecord<>("abc_foo", 0, 1, null, "b"); + ConsumerRecord record3 = new ConsumerRecord<>("abc_foo", 1, 0, null, "c"); + ConsumerRecord record4 = new ConsumerRecord<>("def_foo", 1, 0, null, "d"); + ConsumerRecord record5 = new ConsumerRecord<>("def_foo", 0, 0, null, "e"); + Arrays.asList(record1, record2, record3, record4, record5) + .forEach(consumer::addRecord); + + Message received = source.receive(); + assertThat(received).isNotNull(); + assertThat(received.getHeaders().get(KafkaHeaders.RAW_DATA)).isInstanceOf(ConsumerRecord.class); + assertThat(received.getHeaders().get(KafkaHeaders.RAW_DATA)).isEqualTo(record1); + received = source.receive(); + assertThat(received).isNotNull(); + assertThat(received.getHeaders().get(KafkaHeaders.RAW_DATA)).isInstanceOf(ConsumerRecord.class); + assertThat(received.getHeaders().get(KafkaHeaders.RAW_DATA)).isEqualTo(record2); + received = source.receive(); + assertThat(received).isNotNull(); + assertThat(received.getHeaders().get(KafkaHeaders.RAW_DATA)).isInstanceOf(ConsumerRecord.class); + assertThat(received.getHeaders().get(KafkaHeaders.RAW_DATA)).isEqualTo(record3); + received = source.receive(); + assertThat(received).isNotNull(); + assertThat(received.getHeaders().get(KafkaHeaders.RAW_DATA)).isInstanceOf(ConsumerRecord.class); + assertThat(received.getHeaders().get(KafkaHeaders.RAW_DATA)).isEqualTo(record4); + received = source.receive(); + assertThat(received).isNotNull(); + assertThat(received.getHeaders().get(KafkaHeaders.RAW_DATA)).isInstanceOf(ConsumerRecord.class); + assertThat(received.getHeaders().get(KafkaHeaders.RAW_DATA)).isEqualTo(record5); + + source.destroy(); + } + + @SuppressWarnings({ "rawtypes", "unchecked" }) + @Test + public void testStaticPartitionAssignment() { + MockConsumer consumer = spy(new MockConsumer<>(OffsetResetStrategy.EARLIEST)); + + TopicPartition beginning = new TopicPartition("foo", 0); + TopicPartition end = new TopicPartition("foo", 1); + TopicPartition timestamp = new TopicPartition("foo", 2); + TopicPartition negativeOffset = new TopicPartition("foo", 3); + TopicPartition negativeRelativeToCurrent = new TopicPartition("foo", 4); + TopicPartition positiveRelativeToCurrent = new TopicPartition("foo", 5); + + List topicPartitions = Arrays.asList(beginning, end, timestamp, + negativeOffset, negativeRelativeToCurrent, positiveRelativeToCurrent); + + Map beginningOffsets = topicPartitions.stream().collect(Collectors + .toMap(Function.identity(), tp -> 0L)); + consumer.updateBeginningOffsets(beginningOffsets); + Map endOffsets = topicPartitions.stream().collect(Collectors + .toMap(Function.identity(), tp -> 3L)); + consumer.updateEndOffsets(endOffsets); + consumer.assign(topicPartitions); + consumer.seek(timestamp, 1); + consumer.seek(negativeRelativeToCurrent, 3); + consumer.seek(positiveRelativeToCurrent, 1); + + ConsumerFactory consumerFactory = mock(ConsumerFactory.class); + willReturn(Collections.singletonMap(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, 1)).given(consumerFactory) + .getConfigurationProperties(); + given(consumerFactory.createConsumer(isNull(), anyString(), isNull())).willReturn(consumer); + + TopicPartitionOffset beginningTpo = new TopicPartitionOffset(beginning, null, TopicPartitionOffset.SeekPosition.BEGINNING); + TopicPartitionOffset endTpo = new TopicPartitionOffset(end, null, TopicPartitionOffset.SeekPosition.END); + TopicPartitionOffset timestampTpo = new TopicPartitionOffset(timestamp, null, TopicPartitionOffset.SeekPosition.TIMESTAMP); + TopicPartitionOffset negativeOffsetTpo = new TopicPartitionOffset(negativeOffset, -1L, null); + TopicPartitionOffset negativeRelativeToCurrentTpo = new TopicPartitionOffset(negativeRelativeToCurrent.topic(), + negativeRelativeToCurrent.partition(), -1L, true); + TopicPartitionOffset positiveRelativeToCurrentTpo = new TopicPartitionOffset(positiveRelativeToCurrent.topic(), + positiveRelativeToCurrent.partition(), 1L, true); + KafkaMessageSource source = new KafkaMessageSource(consumerFactory, + new ConsumerProperties(beginningTpo, endTpo, timestampTpo, + negativeOffsetTpo, negativeRelativeToCurrentTpo, positiveRelativeToCurrentTpo)); + source.setRawMessageHeader(true); + + source.createConsumer(); + + ConsumerRecord p0r0 = new ConsumerRecord<>("foo", 0, 0, null, "p0r0"); + ConsumerRecord p0r1 = new ConsumerRecord<>("foo", 0, 1, null, "p0r1"); + ConsumerRecord p0r2 = new ConsumerRecord<>("foo", 0, 2, null, "p0r2"); + ConsumerRecord p0r3 = new ConsumerRecord<>("foo", 0, 3, null, "p0r3"); + ConsumerRecord p1r0 = new ConsumerRecord<>("foo", 1, 0, null, "p1r0"); + ConsumerRecord p1r1 = new ConsumerRecord<>("foo", 1, 1, null, "p1r1"); + ConsumerRecord p1r2 = new ConsumerRecord<>("foo", 1, 2, null, "p1r2"); + ConsumerRecord p1r3 = new ConsumerRecord<>("foo", 1, 3, null, "p1r3"); + ConsumerRecord p2r0 = new ConsumerRecord<>("foo", 2, 0, null, "p2r0"); + ConsumerRecord p2r1 = new ConsumerRecord<>("foo", 2, 1, null, "p2r1"); + ConsumerRecord p2r2 = new ConsumerRecord<>("foo", 2, 2, null, "p2r2"); + ConsumerRecord p2r3 = new ConsumerRecord<>("foo", 2, 3, null, "p2r3"); + ConsumerRecord p3r0 = new ConsumerRecord<>("foo", 3, 0, null, "p3r0"); + ConsumerRecord p3r1 = new ConsumerRecord<>("foo", 3, 1, null, "p3r1"); + ConsumerRecord p3r2 = new ConsumerRecord<>("foo", 3, 2, null, "p3r2"); + ConsumerRecord p3r3 = new ConsumerRecord<>("foo", 3, 3, null, "p3r3"); + ConsumerRecord p4r0 = new ConsumerRecord<>("foo", 4, 0, null, "p4r0"); + ConsumerRecord p4r1 = new ConsumerRecord<>("foo", 4, 1, null, "p4r1"); + ConsumerRecord p4r2 = new ConsumerRecord<>("foo", 4, 2, null, "p4r2"); + ConsumerRecord p4r3 = new ConsumerRecord<>("foo", 4, 3, null, "p4r3"); + ConsumerRecord p5r0 = new ConsumerRecord<>("foo", 5, 0, null, "p5r0"); + ConsumerRecord p5r1 = new ConsumerRecord<>("foo", 5, 1, null, "p5r1"); + ConsumerRecord p5r2 = new ConsumerRecord<>("foo", 5, 2, null, "p5r2"); + ConsumerRecord p5r3 = new ConsumerRecord<>("foo", 5, 3, null, "p5r3"); + Arrays.asList(p0r0, p0r1, p0r2, p0r3, p1r0, p1r1, p1r2, p1r3, p2r0, p2r1, p2r2, p2r3, p3r0, p3r1, p3r2, p3r3, + p4r0, p4r1, p4r2, p4r3, p5r0, p5r1, p5r2, p5r3) + .forEach(consumer::addRecord); + + + Message message; + Set expected = Stream.of( + p0r0, p0r1, p0r2, p0r3, // Seek to beginning + p1r3, // Seek to end + p2r1, p2r2, p2r3, // Null offset and SeekPosition.TIMESTAMP results in no change in position + p3r3, // Negative offset ends up in seek to end + p4r2, p4r3, // Negative offset with relative to current(3) + p5r2, p5r3 // Positive offset with relative to current(1) + ).map(ConsumerRecord::value).collect(Collectors.toSet()); + Set received = new HashSet<>(); + while ((message = source.receive()) != null) { + received.add(message.getHeaders().get(KafkaHeaders.RAW_DATA, ConsumerRecord.class).value()); + } + + assertThat(received).isEqualTo(expected); + + source.pause(); + source.receive(); + source.resume(); + source.receive(); + source.destroy(); + + InOrder inOrder = inOrder(consumer); + inOrder.verify(consumer).assign(anyCollection()); + inOrder.verify(consumer).seekToBeginning(anyCollection()); + inOrder.verify(consumer, times(2)).seekToEnd(anyCollection()); + inOrder.verify(consumer).position(negativeRelativeToCurrent); + inOrder.verify(consumer).seek(negativeRelativeToCurrent, 2L); + inOrder.verify(consumer).position(positiveRelativeToCurrent); + inOrder.verify(consumer).seek(positiveRelativeToCurrent, 2L); + inOrder.verify(consumer).poll(any(Duration.class)); + inOrder.verify(consumer).poll(any(Duration.class)); + inOrder.verify(consumer).pause(anyCollection()); + inOrder.verify(consumer).poll(any(Duration.class)); + inOrder.verify(consumer).resume(anyCollection()); + inOrder.verify(consumer).poll(any(Duration.class)); + inOrder.verify(consumer).close(); + inOrder.verify(consumer).close(anyLong(), any(TimeUnit.class)); + inOrder.verifyNoMoreInteractions(); + } }