raise KafkaTimeoutException when flush times out

Author: Andrew Kowalik

kafka/producer/kafka.py

@@ -548,6 +548,10 @@ def flush(self, timeout=None):
 
         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()

kafka/producer/record_accumulator.py

@@ -522,8 +522,11 @@ def await_flush_completion(self, timeout=None):
             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: