Add KafkaItemReader

Resolves BATCH-2764

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/kafka/KafkaItemReader.java

@@ -0,0 +1,188 @@
+/*
+ * Copyright 2018 the original author or authors.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *          http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+package org.springframework.batch.item.kafka;
+
+import java.time.Duration;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Supplier;
+import java.util.stream.Collectors;
+
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.TopicPartition;
+import org.springframework.batch.item.ExecutionContext;
+import org.springframework.batch.item.ItemReader;
+import org.springframework.batch.item.ItemStreamException;
+import org.springframework.batch.item.kafka.support.AutoCommitOffsetsProvider;
+import org.springframework.batch.item.support.AbstractItemCountingItemStreamItemReader;
+import org.springframework.beans.factory.InitializingBean;
+import org.springframework.kafka.core.ConsumerFactory;
+import org.springframework.util.Assert;
+import org.springframework.util.ClassUtils;
+
+/**
+ * <p>
+ * An {@link ItemReader} implementation for Apache Kafka.
+ * </p>
+ *
+ * @author Mathieu Ouellet
+ * @since 4.2
+ *
+ */
+public class KafkaItemReader<K, V> extends AbstractItemCountingItemStreamItemReader<V> implements InitializingBean {
+
+	private static final String TOPIC_PARTITION_OFFSET = "topic.partition.offset";
+
+	private static final long DEFAULT_POLL_TIMEOUT = 50L;
+
+	private static final long MIN_ASSIGN_TIMEOUT = 2000L;
+
+	private final Supplier<Duration> assignTimeoutProvider = () -> Duration
+			.ofMillis(Math.max(this.pollTimeout.toMillis() * 20, MIN_ASSIGN_TIMEOUT));
+
+	private Duration pollTimeout = Duration.ofMillis(DEFAULT_POLL_TIMEOUT);
+
+	private List<TopicPartition> topicPartitions;
+
+	private List<String> topics;
+
+	private ConsumerFactory<K, V> consumerFactory;
+
+	private Consumer<K, V> consumer;
+
+	private OffsetsProvider offsetsProvider;
+
+	private AtomicBoolean assigned = new AtomicBoolean(false);
+
+	private Map<TopicPartition, Long> offsets;
+
+	private Iterator<ConsumerRecord<K, V>> records;
+
+	public KafkaItemReader() {
+		super();
+		setName(ClassUtils.getShortName(KafkaItemReader.class));
+	}
+
+	public void setPollTimeout(long pollTimeout) {
+		Assert.isTrue(pollTimeout >= 0, "'pollTimeout' must no be negative.");
+		this.pollTimeout = Duration.ofMillis(pollTimeout);
+	}
+
+	public void setTopicPartitions(List<TopicPartition> topicPartitions) {
+		this.topicPartitions = topicPartitions;
+	}
+
+	public void setTopics(List<String> topics) {
+		this.topics = topics;
+	}
+
+	public void setConsumerFactory(ConsumerFactory<K, V> consumerFactory) {
+		this.consumerFactory = consumerFactory;
+	}
+
+	public void setOffsetsProvider(OffsetsProvider offsetsProvider) {
+		this.offsetsProvider = offsetsProvider;
+	}
+
+	@Override
+	public void afterPropertiesSet() throws Exception {
+		Assert.state(topicPartitions != null || topics != null, "Either 'topicPartitions' or 'topics' must be provided.");
+		Assert.state(topicPartitions == null || topics == null, "Both 'topicPartitions' and 'topics' cannot be specified together.");
+		Assert.notNull(consumerFactory, "'consumerFactory' must not be null.");
+		Assert.notNull(offsetsProvider, "'offsetsProvider' must not be null.");
+		if  (consumerFactory.isAutoCommit()) {
+			Assert.state(offsetsProvider instanceof AutoCommitOffsetsProvider, "'AutoCommitOffsetsProvider' must be used if 'consumerFactory' is set to auto commit.");
+		}
+	}
+
+	@Override
+	@SuppressWarnings("unchecked")
+	public void open(ExecutionContext executionContext) throws ItemStreamException {
+		super.open(executionContext);
+		try {
+			if (isSaveState() && executionContext.containsKey(TOPIC_PARTITION_OFFSET)) {
+				offsets = (Map<TopicPartition, Long>) executionContext.get(TOPIC_PARTITION_OFFSET);
+			}
+			else {
+				offsets = offsetsProvider.get(topicPartitions);
+			}
+
+			if (offsets != null && !offsets.isEmpty()) {
+				offsets.forEach(consumer::seek);
+			}
+		}
+		catch (Exception e) {
+			throw new ItemStreamException("Failed to initialize the reader", e);
+		}
+	}
+
+	@Override
+	protected void doOpen() throws Exception {
+		consumer = consumerFactory.createConsumer();
+		if (topics != null) {
+			topicPartitions = topics.stream()
+					.flatMap(topic -> consumer.partitionsFor(topic).stream())
+					.map(partitionInfo -> new TopicPartition(partitionInfo.topic(), partitionInfo.partition()))
+					.collect(Collectors.toList());
+		}
+		consumer.assign(topicPartitions);
+		offsetsProvider.setConsumer(consumer);
+	}
+
+	@Override
+	protected void jumpToItem(int itemIndex) throws Exception {
+	}
+
+	@Override
+	protected V doRead() throws Exception {
+		if (records == null || !records.hasNext()) {
+			records = doPoll();
+		}
+		if (records.hasNext()) {
+			ConsumerRecord<K, V> record = records.next();
+			offsets.put(new TopicPartition(record.topic(), record.partition()), record.offset());
+			return record.value();
+		}
+		return null;
+	}
+
+	protected Iterator<ConsumerRecord<K, V>> doPoll() {
+		return consumer.poll(assigned.getAndSet(true) ? pollTimeout : assignTimeoutProvider.get()).iterator();
+	}
+
+	@Override
+	public void update(ExecutionContext executionContext) throws ItemStreamException {
+		super.update(executionContext);
+		if (isSaveState()) {
+			Assert.notNull(executionContext, "ExecutionContext must not be null");
+			executionContext.put(TOPIC_PARTITION_OFFSET, offsets);
+		}
+	}
+
+	@Override
+	protected void doClose() throws Exception {
+		records = null;
+		offsets = null;
+		assigned.set(false);
+		if (consumer != null) {
+			consumer.close();
+		}
+	}
+}

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/kafka/OffsetsProvider.java

@@ -0,0 +1,37 @@
+package org.springframework.batch.item.kafka;
+
+import java.util.List;
+import java.util.Map;
+
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.common.TopicPartition;
+
+/**
+ * A convenient strategy for providing initial externally stored offsets for the {@link KafkaItemReader} to
+ * seeks to.
+ *
+ * @author Mathieu Ouellet
+ * @since 4.2
+ */
+public interface OffsetsProvider {
+
+	/**
+	 * <p>
+	 * Map of assigned topic-partitions and the offset, or read position, at which the {@link KafkaItemReader} should
+	 * start.
+	 * </p>
+	 *
+	 * @param topicPartitions list of assigned topic-partitions to get offset for
+	 * @return map of offset by topic-partition
+	 */
+	Map<TopicPartition, Long> get(List<TopicPartition> topicPartitions);
+
+	/**
+	 * <p>
+	 * Inject a {@link Consumer} that can be used to fetch internally stored offsets and/or topic-partition assignment.
+	 * </p>
+	 *
+	 * @param consumer the {@link Consumer} to set
+	 */
+	void setConsumer(Consumer<?, ?> consumer);
+}

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/kafka/builder/KafkaItemReaderBuilder.java

@@ -0,0 +1,175 @@
+/*
+ * Copyright 2018 the original author or authors.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *          http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+package org.springframework.batch.item.kafka.builder;
+
+import java.util.List;
+
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.common.TopicPartition;
+import org.springframework.batch.item.kafka.KafkaItemReader;
+import org.springframework.batch.item.kafka.OffsetsProvider;
+import org.springframework.batch.item.kafka.support.AutoCommitOffsetsProvider;
+import org.springframework.kafka.core.ConsumerFactory;
+import org.springframework.util.Assert;
+
+/**
+ * A builder implementation for the {@link KafkaItemReader}
+ *
+ * @author Mathieu Ouellet
+ * @since 4.2
+ * @see KafkaItemReader
+ */
+public class KafkaItemReaderBuilder<K, V> {
+
+	private ConsumerFactory<K, V> consumerFactory;
+
+	private List<TopicPartition> topicPartitions;
+
+	private List<String> topics;
+
+	private OffsetsProvider offsetsProvider = new AutoCommitOffsetsProvider();
+
+	private long pollTimeout = 50L;
+
+	private boolean saveState = true;
+
+	private String name;
+
+	private int maxItemCount = Integer.MAX_VALUE;
+
+	/**
+	 * Configure if the state of the {@link org.springframework.batch.item.ItemStreamSupport} should be persisted within
+	 * the {@link org.springframework.batch.item.ExecutionContext} for restart purposes.
+	 *
+	 * @param saveState defaults to true
+	 * @return The current instance of the builder.
+	 */
+	public KafkaItemReaderBuilder<K, V> saveState(boolean saveState) {
+		this.saveState = saveState;
+		return this;
+	}
+
+	/**
+	 * The name used to calculate the key within the {@link org.springframework.batch.item.ExecutionContext}. Required
+	 * if {@link #saveState(boolean)} is set to true.
+	 *
+	 * @param name name of the reader instance
+	 * @return The current instance of the builder.
+	 * @see org.springframework.batch.item.ItemStreamSupport#setName(String)
+	 */
+	public KafkaItemReaderBuilder<K, V> name(String name) {
+		this.name = name;
+		return this;
+	}
+
+	/**
+	 * Configure the max number of items to be read.
+	 *
+	 * @param maxItemCount the max items to be read
+	 * @return The current instance of the builder.
+	 * @see org.springframework.batch.item.support.AbstractItemCountingItemStreamItemReader#setMaxItemCount(int)
+	 */
+	public KafkaItemReaderBuilder<K, V> maxItemCount(int maxItemCount) {
+		this.maxItemCount = maxItemCount;
+		return this;
+	}
+
+	/**
+	 * The {@link ConsumerFactory} implementation to produce a new {@link Consumer} instance for the reader.
+	 *
+	 * @param consumerFactory
+	 * @return The current instance of the builder.
+	 * @see KafkaItemReader#setConsumerFactory(ConsumerFactory)
+	 */
+	public KafkaItemReaderBuilder<K, V> consumerFactory(ConsumerFactory<K, V> consumerFactory) {
+		this.consumerFactory = consumerFactory;
+		return this;
+	}
+
+	/**
+	 * A list of {@link TopicPartition}s to manually assign the consumer.
+	 *
+	 * @param topicPartitions list of partitions to assign the consumer
+	 * @return The current instance of the builder.
+	 * @see KafkaItemReader#setTopicPartitions(List)
+	 */
+	public KafkaItemReaderBuilder<K, V> topicPartitions(List<TopicPartition> topicPartitions) {
+		this.topicPartitions = topicPartitions;
+		return this;
+	}
+
+	/**
+	 * A list of topics to manually assign the consumer.
+	 *
+	 * @param topics list of topics to assign the consumer
+	 * @return The current instance of the builder.
+	 * @see KafkaItemReader#setTopics(List)
+	 */
+	public KafkaItemReaderBuilder<K, V> topics(List<String> topics) {
+		this.topics = topics;
+		return this;
+	}
+
+	/**
+	 * The {@link OffsetsProvider} implementation to provide initial offsets.
+	 *
+	 * @param offsetsProvider
+	 * @return The current instance of the builder.
+	 * @see KafkaItemReader#setOffsetsProvider(OffsetsProvider)
+	 */
+	public KafkaItemReaderBuilder<K,V> offsetsProvider(OffsetsProvider offsetsProvider)  {
+		this.offsetsProvider = offsetsProvider;
+		return this;
+	}
+
+	/**
+	 * Set the pollTimeout for the poll() operations.
+	 *
+	 * @param pollTimeout default to 50ms
+	 * @return The current instance of the builder.
+	 * @see KafkaItemReader#setPollTimeout(long)
+	 */
+	public KafkaItemReaderBuilder<K, V> pollTimeout(long pollTimeout) {
+		this.pollTimeout = pollTimeout;
+		return this;
+	}
+
+	public KafkaItemReader<K, V> build() {
+		if (this.saveState) {
+			Assert.hasText(this.name, "A name is required when saveState is set to true");
+		}
+		Assert.state(this.topicPartitions != null || this.topics != null, "Either 'topicPartitions' or 'topics' must be provided.");
+		Assert.state(this.topicPartitions == null || this.topics == null, "Both 'topicPartitions' and 'topics' cannot be specified together.");
+		Assert.isTrue(this.pollTimeout >= 0, "pollTimeout must not be negative.");
+		Assert.notNull(this.consumerFactory, "'consumerFactory' must not be null.");
+		Assert.notNull(this.offsetsProvider, "'offsetsProvider' must not be null.");
+		if  (this.consumerFactory.isAutoCommit()) {
+			Assert.state(this.offsetsProvider instanceof AutoCommitOffsetsProvider, "'AutoCommitOffsetsProvider' must be used if 'consumerFactory' is set to auto commit.");
+		}
+
+		KafkaItemReader<K, V> reader = new KafkaItemReader<>();
+		reader.setConsumerFactory(this.consumerFactory);
+		reader.setTopicPartitions(this.topicPartitions);
+		reader.setTopics(this.topics);
+		reader.setOffsetsProvider(this.offsetsProvider);
+		reader.setPollTimeout(this.pollTimeout);
+		reader.setSaveState(this.saveState);
+		reader.setName(this.name);
+		reader.setMaxItemCount(this.maxItemCount);
+		return reader;
+	}
+}

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/kafka/support/AutoCommitOffsetsProvider.java

@@ -0,0 +1,32 @@
+package org.springframework.batch.item.kafka.support;
+
+import java.util.List;
+import java.util.Map;
+
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.common.TopicPartition;
+import org.springframework.batch.item.kafka.OffsetsProvider;
+
+/**
+ * Noop implementation of {@link OffsetsProvider} to use with automatic offset committing. This is the only
+ * implementation allowed when the KafkaConsumer has 'enable.auto.commit' set to true. See 'auto.offset.reset' config
+ * for initial offset strategy options.
+ *
+ * @author Mathieu Ouellet
+ * @since 4.2
+ */
+public class AutoCommitOffsetsProvider implements OffsetsProvider {
+
+	/**
+	 * @return null, preventing a call to
+	 * {@link org.apache.kafka.clients.consumer.Consumer#seek(org.apache.kafka.common.TopicPartition, long)}
+	 */
+	@Override
+	public Map<TopicPartition, Long> get(List<TopicPartition> topicPartitions) {
+		return null;
+	}
+
+	@Override
+	public void setConsumer(Consumer<?, ?> consumer) {
+	}
+}

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/kafka/support/BeginningOffsetsProvider.java

@@ -0,0 +1,34 @@
+package org.springframework.batch.item.kafka.support;
+
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.common.TopicPartition;
+import org.springframework.batch.item.kafka.OffsetsProvider;
+
+/**
+ * <p>
+ * Implementation of {@link OffsetsProvider} that returns the earliest offsets for the given topic-partitions.
+ * Equivalent of using 'auto.offset.reset' set to 'earliest' with 'enable.auto.commit' set to false.
+ * </p>
+ *
+ * @author Mathieu Ouellet
+ * @see org.apache.kafka.clients.consumer.KafkaConsumer#beginningOffsets(Collection)
+ * @since 4.2
+ */
+public class BeginningOffsetsProvider implements OffsetsProvider {
+
+	private Consumer<?, ?> consumer;
+
+	@Override
+	public Map<TopicPartition, Long> get(List<TopicPartition> topicPartitions) {
+		return this.consumer.beginningOffsets(topicPartitions);
+	}
+
+	@Override
+	public void setConsumer(Consumer<?, ?> consumer) {
+		this.consumer = consumer;
+	}
+}

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/kafka/support/SimpleOffsetsProvider.java

@@ -0,0 +1,34 @@
+package org.springframework.batch.item.kafka.support;
+
+import java.util.List;
+import java.util.Map;
+
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.common.TopicPartition;
+import org.springframework.batch.item.kafka.OffsetsProvider;
+
+/**
+ * <p>
+ * Implementation of {@link OffsetsProvider} that returns default or provided offsets for the given topic-partitions.
+ * </p>
+ *
+ * @author Mathieu Ouellet
+ * @since 4.2
+ */
+public class SimpleOffsetsProvider implements OffsetsProvider {
+
+	private Map<TopicPartition, Long> offsets;
+
+	@Override
+	public Map<TopicPartition, Long> get(List<TopicPartition> topicPartitions) {
+		return this.offsets;
+	}
+
+	@Override
+	public void setConsumer(Consumer<?, ?> consumer) {
+	}
+
+	public void setOffsets(Map<TopicPartition, Long> offsets) {
+		this.offsets = offsets;
+	}
+}

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/kafka/support/TimestampOffsetsProvider.java

@@ -0,0 +1,45 @@
+package org.springframework.batch.item.kafka.support;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.common.TopicPartition;
+import org.springframework.batch.item.kafka.OffsetsProvider;
+
+/**
+ * <p>
+ * Implementation of {@link OffsetsProvider} that returns offsets for the given topic-partitions by timestamp.
+ * </p>
+ *
+ * @author Mathieu Ouellet
+ * @see org.apache.kafka.clients.consumer.KafkaConsumer#offsetsForTimes(java.util.Map)
+ * @since 4.2
+ */
+public class TimestampOffsetsProvider implements OffsetsProvider {
+
+	private final Long timestampToSearch;
+
+	private Consumer<?, ?> consumer;
+
+	public TimestampOffsetsProvider(Long timestampToSearch) {
+		this.timestampToSearch = timestampToSearch;
+	}
+
+	@Override
+	public Map<TopicPartition, Long> get(List<TopicPartition> topicPartitions) {
+		Map<TopicPartition, Long> timestampsToSearch = topicPartitions.stream()
+				.collect(Collectors.toMap(Function.identity(), topicPartition -> timestampToSearch));
+		return consumer.offsetsForTimes(timestampsToSearch).entrySet().stream()
+				.filter(entry -> entry.getValue() != null)
+				.collect(Collectors.toMap(Entry::getKey, entry -> entry.getValue().offset()));
+	}
+
+	@Override
+	public void setConsumer(Consumer<?, ?> consumer) {
+		this.consumer = consumer;
+	}
+}

spring-batch-infrastructure/src/test/java/org/springframework/batch/item/kafka/KafkaItemReaderTests.java

@@ -0,0 +1,167 @@
+package org.springframework.batch.item.kafka;
+
+import java.time.Duration;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.common.TopicPartition;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.springframework.batch.item.ExecutionContext;
+import org.springframework.kafka.core.ConsumerFactory;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.singletonList;
+
+import static org.hamcrest.Matchers.is;
+import static org.junit.Assert.assertThat;
+import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.BDDMockito.given;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+/**
+ * @author Mathieu Ouellet
+ */
+public class KafkaItemReaderTests {
+
+	private static final TopicPartition TOPIC_PARTITION = new TopicPartition("topic", 0);
+
+	@Mock
+	private ConsumerFactory<String, String> consumerFactory;
+
+	@Mock
+	private Consumer<String, String> consumer;
+
+	@Mock
+	private OffsetsProvider offsetsProvider;
+
+	private KafkaItemReader<String, String> reader;
+
+	@Before
+	public void setUp() throws Exception {
+		MockitoAnnotations.initMocks(this);
+		Map<String, Object> config = new HashMap<>();
+		config.put("max.poll.records", 2);
+		config.put("enable.auto.commit", false);
+		when(consumerFactory.getConfigurationProperties()).thenReturn(config);
+		when(consumerFactory.createConsumer()).thenReturn(consumer);
+		reader = new KafkaItemReader<>();
+		reader.setConsumerFactory(consumerFactory);
+		reader.setOffsetsProvider(offsetsProvider);
+		reader.setTopicPartitions(singletonList(TOPIC_PARTITION));
+		reader.setSaveState(true);
+		reader.setPollTimeout(50L);
+		reader.afterPropertiesSet();
+	}
+
+	@Test
+	public void testAfterPropertiesSet() throws Exception {
+		reader = new KafkaItemReader<>();
+
+		try {
+			reader.afterPropertiesSet();
+			fail("Expected exception was not thrown");
+		}
+		catch (IllegalStateException ignore) {
+		}
+
+		reader.setTopicPartitions(singletonList(new TopicPartition("topic", 0)));
+		reader.setTopics(singletonList("topic"));
+		try {
+			reader.afterPropertiesSet();
+			fail("Expected exception was not thrown");
+		}
+		catch (IllegalStateException ignore) {
+		}
+
+		reader.setTopics(null);
+		try {
+			reader.afterPropertiesSet();
+			fail("Expected exception was not thrown");
+		}
+		catch (IllegalArgumentException ignore) {
+		}
+
+		reader.setConsumerFactory(consumerFactory);
+		try {
+			reader.afterPropertiesSet();
+			fail("Expected exception was not thrown");
+		}
+		catch (IllegalArgumentException ignore) {
+		}
+
+		reader.setOffsetsProvider(offsetsProvider);
+		reader.afterPropertiesSet();
+	}
+
+	@Test
+	public void testAssignTopicPartitions() {
+		reader.open(new ExecutionContext());
+		verify(consumer).assign(singletonList(TOPIC_PARTITION));
+	}
+
+	@Test
+	public void testRead() throws Exception {
+		Map<TopicPartition, List<ConsumerRecord<String, String>>> records = new HashMap<>();
+		records.put(TOPIC_PARTITION, singletonList(
+				new ConsumerRecord<>(TOPIC_PARTITION.topic(), TOPIC_PARTITION.partition(), 0L, "key0", "val0")));
+		when(consumer.poll(any())).thenReturn(new ConsumerRecords<>(records));
+
+		reader.open(new ExecutionContext());
+		String read = reader.read();
+		assertThat(read, is("val0"));
+	}
+
+	@Test
+	public void testPollRecords() throws Exception {
+		Map<TopicPartition, List<ConsumerRecord<String, String>>> firstPoll = new HashMap<>();
+		firstPoll.put(TOPIC_PARTITION, asList(new ConsumerRecord<>("topic", 0, 0L, "key0", "val0"),
+				new ConsumerRecord<>("topic", 0, 1L, "key1", "val1")));
+		when(consumer.poll(Duration.ofMillis(2000L))).thenReturn(new ConsumerRecords<>(firstPoll));
+
+		Map<TopicPartition, List<ConsumerRecord<String, String>>> secondPoll = new HashMap<>();
+		secondPoll.put(TOPIC_PARTITION, singletonList(new ConsumerRecord<>("topic", 0, 2L, "key2", "val2")));
+		when(consumer.poll(Duration.ofMillis(50L))).thenReturn(new ConsumerRecords<>(secondPoll));
+
+		reader.open(new ExecutionContext());
+
+		String read = reader.read();
+		assertThat(read, is("val0"));
+
+		read = reader.read();
+		assertThat(read, is("val1"));
+
+		read = reader.read();
+		assertThat(read, is("val2"));
+	}
+
+	@Test
+	public void testSeekOnSavedState() {
+		long offset = 100L;
+		Map<TopicPartition, Long> offsets = new HashMap<>();
+		offsets.put(TOPIC_PARTITION, offset);
+		ExecutionContext executionContext = new ExecutionContext();
+		executionContext.put("topic.partition.offset", offsets);
+		reader.open(executionContext);
+		verify(consumer).seek(TOPIC_PARTITION, offset);
+	}
+
+	@Test
+	public void testSeekToProvidedOffsets() {
+		long offset = 100L;
+		Map<TopicPartition, Long> offsets = new HashMap<>();
+		offsets.put(TOPIC_PARTITION, offset);
+		given(offsetsProvider.get(singletonList(TOPIC_PARTITION))).willReturn(offsets);
+		reader.open(new ExecutionContext());
+		verify(consumer).seek(TOPIC_PARTITION, offset);
+	}
+
+}

spring-batch-infrastructure/src/test/java/org/springframework/batch/item/kafka/builder/KafkaItemReaderBuilderTests.java

@@ -0,0 +1,143 @@
+package org.springframework.batch.item.kafka.builder;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.mockito.Mockito.when;
+
+import java.time.Duration;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.apache.kafka.common.TopicPartition;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.springframework.batch.item.kafka.KafkaItemReader;
+import org.springframework.batch.item.kafka.OffsetsProvider;
+import org.springframework.batch.item.kafka.support.BeginningOffsetsProvider;
+import org.springframework.kafka.core.ConsumerFactory;
+import org.springframework.test.util.ReflectionTestUtils;
+
+/**
+ * @author Mathieu Ouellet
+ */
+public class KafkaItemReaderBuilderTests {
+
+	@Rule
+	public ExpectedException thrown = ExpectedException.none();
+
+	@Mock
+	private ConsumerFactory<Object, Object> consumerFactory;
+	@Mock
+	private OffsetsProvider offsetsProvider;
+
+	private List<TopicPartition> topicPartitions = Collections.singletonList(new TopicPartition("topic", 0));
+
+	@Before
+	public void setUp() {
+		MockitoAnnotations.initMocks(this);
+		Map<String, Object> config = new HashMap<>();
+		config.put("max.poll.records", 2);
+		config.put("enable.auto.commit", false);
+		when(consumerFactory.getConfigurationProperties()).thenReturn(config);
+	}
+
+	@Test
+	public void testNullConsumerFactory() {
+		this.thrown.expect(IllegalArgumentException.class);
+		this.thrown.expectMessage("'consumerFactory' must not be null.");
+
+		new KafkaItemReaderBuilder<>()
+				.name("kafkaItemReader")
+				.topicPartitions(topicPartitions)
+				.consumerFactory(null)
+				.offsetsProvider(offsetsProvider)
+				.build();
+	}
+
+	@Test
+	public void testNullTopicsAndTopicPartitions() {
+		this.thrown.expect(IllegalStateException.class);
+		this.thrown.expectMessage("Either 'topicPartitions' or 'topics' must be provided.");
+
+		new KafkaItemReaderBuilder<>()
+				.name("kafkaItemReader")
+				.topicPartitions(null)
+				.consumerFactory(consumerFactory)
+				.offsetsProvider(offsetsProvider)
+				.build();
+	}
+
+	@Test
+	public void testPollTimeoutNegative() {
+		this.thrown.expect(IllegalArgumentException.class);
+		this.thrown.expectMessage("pollTimeout must not be negative.");
+
+		new KafkaItemReaderBuilder<>()
+				.name("kafkaItemReader")
+				.consumerFactory(consumerFactory)
+				.topicPartitions(topicPartitions)
+				.offsetsProvider(offsetsProvider)
+				.pollTimeout(-1)
+				.build();
+	}
+
+	@Test
+	public void testNullOffsetsProvider() {
+		this.thrown.expect(IllegalArgumentException.class);
+		this.thrown.expectMessage("'offsetsProvider' must not be null.");
+
+		new KafkaItemReaderBuilder<>()
+				.name("kafkaItemReader")
+				.consumerFactory(consumerFactory)
+				.topicPartitions(topicPartitions)
+				.offsetsProvider(null)
+				.build();
+	}
+
+	@Test
+	public void testOffsetsProviderWithAutoCommitConsumerFactory() {
+		this.thrown.expect(IllegalStateException.class);
+		this.thrown.expectMessage("'AutoCommitOffsetsProvider' must be used if 'consumerFactory' is set to auto commit.");
+
+		when(consumerFactory.isAutoCommit()).thenReturn(true);
+
+		new KafkaItemReaderBuilder<>()
+				.name("kafkaItemReader")
+				.consumerFactory(consumerFactory)
+				.topicPartitions(topicPartitions)
+				.offsetsProvider(new BeginningOffsetsProvider())
+				.build();
+	}
+
+	@Test
+	public void testKafkaItemReaderBuild() {
+		// given
+		boolean saveState = false;
+		long pollTimeout = 100;
+		int maxItemCount = 100;
+
+		// when
+		KafkaItemReader<Object, Object> reader = new KafkaItemReaderBuilder<>()
+				.consumerFactory(consumerFactory)
+				.topicPartitions(topicPartitions)
+				.offsetsProvider(offsetsProvider)
+				.pollTimeout(pollTimeout)
+				.saveState(saveState)
+				.name("kafkaItemReader")
+				.maxItemCount(maxItemCount)
+				.build();
+
+		// then
+		assertEquals(consumerFactory, ReflectionTestUtils.getField(reader, "consumerFactory"));
+		assertEquals(topicPartitions, ReflectionTestUtils.getField(reader, "topicPartitions"));
+		assertEquals(Duration.ofMillis(pollTimeout), ReflectionTestUtils.getField(reader, "pollTimeout"));
+		assertEquals(saveState, ReflectionTestUtils.getField(reader, "saveState"));
+		assertEquals(maxItemCount, ReflectionTestUtils.getField(reader, "maxItemCount"));
+	}
+}

spring-batch-infrastructure/src/test/java/org/springframework/batch/item/kafka/support/BeginningOffsetsProviderTests.java

@@ -0,0 +1,36 @@
+package org.springframework.batch.item.kafka.support;
+
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.common.TopicPartition;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import static java.util.Collections.singletonList;
+
+import static org.mockito.Mockito.verify;
+
+public class BeginningOffsetsProviderTests {
+
+	private static final TopicPartition TOPIC_PARTITION = new TopicPartition("topic", 0);
+
+	@Mock
+	private Consumer<?, ?> consumer;
+
+	private BeginningOffsetsProvider offsetsProvider;
+
+	@Before
+	public void setUp() throws Exception {
+		MockitoAnnotations.initMocks(this);
+		offsetsProvider = new BeginningOffsetsProvider();
+		offsetsProvider.setConsumer(consumer);
+	}
+
+	@Test
+	public void testGetBeginningOffsets() {
+		offsetsProvider.get(singletonList(TOPIC_PARTITION));
+
+		verify(consumer).beginningOffsets(singletonList(TOPIC_PARTITION));
+	}
+}

spring-batch-infrastructure/src/test/java/org/springframework/batch/item/kafka/support/SimpleOffsetsProviderTests.java

@@ -0,0 +1,36 @@
+package org.springframework.batch.item.kafka.support;
+
+import java.util.Map;
+
+import org.apache.kafka.common.TopicPartition;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.MockitoAnnotations;
+
+import static java.util.Collections.singletonList;
+import static java.util.Collections.singletonMap;
+
+import static org.junit.Assert.assertEquals;
+
+public class SimpleOffsetsProviderTests {
+
+	private static final TopicPartition TOPIC_PARTITION = new TopicPartition("topic", 0);
+
+	private Map<TopicPartition, Long> offsets;
+
+	private SimpleOffsetsProvider offsetsProvider;
+
+	@Before
+	public void setUp() throws Exception {
+		MockitoAnnotations.initMocks(this);
+		offsets = singletonMap(TOPIC_PARTITION, 0L);
+		offsetsProvider = new SimpleOffsetsProvider();
+		offsetsProvider.setOffsets(offsets);
+	}
+
+	@Test
+	public void testGetProvidedOffsets() {
+		assertEquals(offsets, offsetsProvider.get(singletonList(TOPIC_PARTITION)));
+	}
+
+}

spring-batch-infrastructure/src/test/java/org/springframework/batch/item/kafka/support/TimestampOffsetsProviderTests.java

@@ -0,0 +1,57 @@
+package org.springframework.batch.item.kafka.support;
+
+import java.util.Map;
+
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.common.TopicPartition;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import static java.util.Collections.singletonList;
+import static java.util.Collections.singletonMap;
+
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class TimestampOffsetsProviderTests {
+
+	private static final TopicPartition TOPIC_PARTITION = new TopicPartition("topic", 0);
+
+	@Mock
+	private Consumer<?, ?> consumer;
+
+	private Long timestampToSearch = 0L;
+
+	private TimestampOffsetsProvider offsetsProvider;
+
+	@Before
+	public void setUp() throws Exception {
+		MockitoAnnotations.initMocks(this);
+
+		offsetsProvider = new TimestampOffsetsProvider(timestampToSearch);
+		offsetsProvider.setConsumer(consumer);
+	}
+
+	@Test
+	public void testFilterNullOffsetsForTimestamp() {
+		// given
+		when(consumer.offsetsForTimes(singletonMap(TOPIC_PARTITION, timestampToSearch)))
+				.thenReturn(singletonMap(TOPIC_PARTITION, null));
+
+		// when
+		Map<TopicPartition, Long> offsets = offsetsProvider.get(singletonList(TOPIC_PARTITION));
+
+		// then
+		assertTrue(offsets.isEmpty());
+	}
+
+	@Test
+	public void testGetOffsetsForTimes() {
+		offsetsProvider.get(singletonList(TOPIC_PARTITION));
+
+		verify(consumer).offsetsForTimes(singletonMap(TOPIC_PARTITION, timestampToSearch));
+	}
+}