Catch more producer exceptions in failover tests

1 files changed, 6 insertions, 3 deletions

diff --git a/test/test_failover_integration.py b/test/test_failover_integration.py
index 58e9463..f03dfd9 100644
--- a/test/test_failover_integration.py
+++ b/test/test_failover_integration.py
@@ -4,7 +4,8 @@ import time
 
 from kafka import SimpleClient, SimpleConsumer, KeyedProducer
 from kafka.errors import (
-    FailedPayloadsError, ConnectionError, RequestTimedOutError)
+    FailedPayloadsError, ConnectionError, RequestTimedOutError,
+    NotLeaderForPartitionError)
 from kafka.producer.base import Producer
 from kafka.structs import TopicPartition
 
@@ -79,7 +80,8 @@ class TestFailover(KafkaIntegrationTestCase):
                 producer.send_messages(topic, partition, b'success')
                 log.debug("success!")
                 recovered = True
-            except (FailedPayloadsError, ConnectionError, RequestTimedOutError):
+            except (FailedPayloadsError, ConnectionError, RequestTimedOutError,
+                    NotLeaderForPartitionError):
                 log.debug("caught exception sending message -- will retry")
                 continue
 
@@ -166,7 +168,8 @@ class TestFailover(KafkaIntegrationTestCase):
                 producer.send_messages(topic, key, msg)
                 if producer.partitioners[topic].partition(key) == 0:
                     recovered = True
-            except (FailedPayloadsError, ConnectionError):
+            except (FailedPayloadsError, ConnectionError, RequestTimedOutError,
+                    NotLeaderForPartitionError):
                 log.debug("caught exception sending message -- will retry")
                 continue