Code review comment: Please add unit tests to fully cover this logic.

Code before review:
public void start(Receiver<V> sparkReceiver) {
this.sparkReceiver = sparkReceiver;
final SerializableFunction<Object[], Void> storeFn =
(input) -> {
if (input == null) {
return null;
}
if (input[0] instanceof ByteBuffer) {
final ByteBuffer byteBuffer = ((ByteBuffer) input[0]).asReadOnlyBuffer();
final byte[] bytes = new byte[byteBuffer.limit()];
byteBuffer.get(bytes);
final V record = SerializationUtils.deserialize(bytes);
recordsQueue.offer(record);
} else if (input[0] instanceof Iterator) {
final Iterator<V> iterator = (Iterator<V>) input[0];
while (iterator.hasNext()) {
V record = (V) iterator.next();
recordsQueue.offer(record);
}
} else if (input[0] instanceof ArrayBuffer) {
final ArrayBuffer<V> arrayBuffer = (ArrayBuffer<V>) input[0];
final Iterator<V> iterator = arrayBuffer.iterator();
while (iterator.hasNext()) {
V record = (V) iterator.next();
recordsQueue.offer(record);
}
} else {
V record = (V) input[0];
recordsQueue.offer(record);
}
return null;
};
try {
new WrappedSupervisor(sparkReceiver, new SparkConf(), storeFn);
} catch (Exception e) {
LOG.error("Can not init Spark Receiver!", e);
throw new IllegalStateException("Spark Receiver was not initialized");
}
((HasOffset) sparkReceiver).setStartOffset(startOffset);
sparkReceiver.supervisor().startReceiver();
try {
TimeUnit.MILLISECONDS.sleep(START_POLL_TIMEOUT_MS);
} catch (InterruptedException e) {
LOG.error("SparkReceiver was interrupted before polling started", e);
throw new IllegalStateException("Spark Receiver was interrupted before polling started");
}
}

Line under review: while (iterator.hasNext()) {

Code after addressing the review:
public void start(Receiver<V> sparkReceiver) {
this.sparkReceiver = sparkReceiver;
final SerializableFunction<Object[], Void> storeFn =
(input) -> {
if (input == null) {
return null;
}
/*
Use only [0] element - data.
The other elements are not needed because they are related to Spark environment options.
*/
Object data = input[0];
if (data instanceof ByteBuffer) {
final ByteBuffer byteBuffer = ((ByteBuffer) data).asReadOnlyBuffer();
final byte[] bytes = new byte[byteBuffer.limit()];
byteBuffer.get(bytes);
final V record = SerializationUtils.deserialize(bytes);
recordsQueue.offer(record);
} else if (data instanceof Iterator) {
final Iterator<V> iterator = (Iterator<V>) data;
while (iterator.hasNext()) {
V record = iterator.next();
recordsQueue.offer(record);
}
} else if (data instanceof ArrayBuffer) {
final ArrayBuffer<V> arrayBuffer = (ArrayBuffer<V>) data;
final Iterator<V> iterator = arrayBuffer.iterator();
while (iterator.hasNext()) {
V record = iterator.next();
recordsQueue.offer(record);
}
} else {
V record = (V) data;
recordsQueue.offer(record);
}
return null;
};
try {
new WrappedSupervisor(sparkReceiver, new SparkConf(), storeFn);
} catch (Exception e) {
LOG.error("Can not init Spark Receiver!", e);
throw new IllegalStateException("Spark Receiver was not initialized");
}
((HasOffset) sparkReceiver).setStartOffset(startOffset);
sparkReceiver.supervisor().startReceiver();
try {
TimeUnit.MILLISECONDS.sleep(START_POLL_TIMEOUT_MS);
} catch (InterruptedException e) {
LOG.error("SparkReceiver was interrupted before polling started", e);
throw new IllegalStateException("Spark Receiver was interrupted before polling started");
}
}