raise KafkaTimeoutException when flush times out

2 files changed, 9 insertions, 2 deletions

diff --git a/kafka/producer/kafka.py b/kafka/producer/kafka.py
index 3abadcc..57155e5 100644
--- a/kafka/producer/kafka.py
+++ b/kafka/producer/kafka.py
@@ -554,6 +554,10 @@ class KafkaProducer(object):
 
         Arguments:
             timeout (float, optional): timeout in seconds to wait for completion.
+            
+        Raises:
+            KafkaTimeoutError: failure to flush buffered records within the 
+                provided timeout 
         """
         log.debug("Flushing accumulated records in producer.")  # trace
         self._accumulator.begin_flush()
diff --git a/kafka/producer/record_accumulator.py b/kafka/producer/record_accumulator.py
index fd081aa..3e97fd7 100644
--- a/kafka/producer/record_accumulator.py
+++ b/kafka/producer/record_accumulator.py
@@ -526,8 +526,11 @@ class RecordAccumulator(object):
             for batch in self._incomplete.all():
                 log.debug('Waiting on produce to %s',
                           batch.produce_future.topic_partition)
-                assert batch.produce_future.wait(timeout=timeout), 'Timeout waiting for future'
-                assert batch.produce_future.is_done, 'Future not done?'
+                if not batch.produce_future.wait(timeout=timeout):
+                    raise Errors.KafkaTimeoutError('Timeout waiting for future')
+                if not batch.produce_future.is_done:
+                    raise Errors.UnknownError('Future not done')
+
                 if batch.produce_future.failed():
                     log.warning(batch.produce_future.exception)
         finally: