diff --git a/kafka/producer/producer_batch.py b/kafka/producer/producer_batch.py index 8be08f575..198a3a0c5 100644 --- a/kafka/producer/producer_batch.py +++ b/kafka/producer/producer_batch.py @@ -180,5 +180,6 @@ def __str__(self): return 'ProducerBatch(topic_partition=%s, record_count=%d)' % ( self.topic_partition, self.records.next_offset()) - - + # for heapq + def __lt__(self, other): + return self.created < other.created diff --git a/test/test_producer_batch.py b/test/test_producer_batch.py index bffa79fcb..7d959cbe9 100644 --- a/test/test_producer_batch.py +++ b/test/test_producer_batch.py @@ -134,3 +134,18 @@ def test_complete_exceptionally_with_null_record_errors(batch): with pytest.raises(AssertionError): _test_complete_exceptionally(batch, record_count, top_level_exception, None) + + +def test_producer_batch_lt(tp, memory_records_builder): + b1 = ProducerBatch(tp, memory_records_builder, now=1) + b2 = ProducerBatch(tp, memory_records_builder, now=2) + + assert b1 < b2 + assert not b1 < b1 + + import heapq + q = [] + heapq.heappush(q, b2) + heapq.heappush(q, b1) + assert q[0] == b1 + assert q[1] == b2