Fix Fetcher.PartitionRecords to handle fetch_offset in the middle of compressed messageset

1 files changed, 23 insertions, 0 deletions

diff --git a/test/test_fetcher.py b/test/test_fetcher.py
index 429071a..aa8e9c3 100644
--- a/test/test_fetcher.py
+++ b/test/test_fetcher.py
@@ -498,3 +498,26 @@ def test__parse_fetched_data__out_of_range(fetcher, topic, mocker):
     partition_record = fetcher._parse_fetched_data(completed_fetch)
     assert partition_record is None
     assert fetcher._subscriptions.assignment[tp].awaiting_reset is True
+
+
+def test_partition_records_offset():
+    """Test that compressed messagesets are handle correctly
+    when fetch offset is in the middle of the message list
+    """
+    batch_start = 120
+    batch_end = 130
+    fetch_offset = 123
+    tp = TopicPartition('foo', 0)
+    messages = [ConsumerRecord(tp.topic, tp.partition, i,
+                               None, None, 'key', 'value', 'checksum', 0, 0)
+                for i in range(batch_start, batch_end)]
+    records = Fetcher.PartitionRecords(fetch_offset, None, messages)
+    assert records.has_more()
+    msgs = records.take(1)
+    assert msgs[0].offset == 123
+    assert records.fetch_offset == 124
+    msgs = records.take(2)
+    assert len(msgs) == 2
+    assert records.has_more()
+    records.discard()
+    assert not records.has_more()