test_batched_simple_producer__triggers_by_message should wait for producer queue to drain before testing messages were sent

1 files changed, 10 insertions, 0 deletions

diff --git a/test/test_producer_integration.py b/test/test_producer_integration.py
index e3f7767..329e4bc 100644
--- a/test/test_producer_integration.py
+++ b/test/test_producer_integration.py
@@ -251,6 +251,16 @@ class TestKafkaProducerIntegration(KafkaIntegrationTestCase):
         # Batch mode is async. No ack
         self.assertEqual(len(resp), 0)
 
+        # Wait until producer has pulled all messages from internal queue
+        # this should signal that the first batch was sent, and the producer
+        # is now waiting for enough messages to batch again (or a timeout)
+        timeout = 5
+        start = time.time()
+        while not producer.queue.empty():
+            if time.time() - start > timeout:
+                self.fail('timeout waiting for producer queue to empty')
+            time.sleep(0.1)
+
         # send messages groups all *msgs in a single call to the same partition
         # so we should see all messages from the first call in one partition
         self.assert_fetch_offset(partitions[0], start_offsets[0], [