catch all errors thrown by _get_leader_for_partition in SimpleClient

zackdever · zackdever · commit 4915942b0723 · 2016-03-17T15:39:28.000-07:00
diff --git a/kafka/client.py b/kafka/client.py
@@ -169,7 +169,8 @@ def _payloads_by_broker(self, payloads):
         for payload in payloads:
             try:
                 leader = self._get_leader_for_partition(payload.topic, payload.partition)
-            except KafkaUnavailableError:
+            except (KafkaUnavailableError, LeaderNotAvailableError,
+                    UnknownTopicOrPartitionError):
                 leader = None
             payloads_by_broker[leader].append(payload)
         return dict(payloads_by_broker)
diff --git a/test/test_client.py b/test/test_client.py
@@ -11,7 +11,7 @@
     BrokerMetadata,
     TopicPartition, KafkaUnavailableError,
     LeaderNotAvailableError, UnknownTopicOrPartitionError,
-    KafkaTimeoutError, ConnectionError
+    KafkaTimeoutError, ConnectionError, FailedPayloadsError
 )
 from kafka.conn import KafkaConnection
 from kafka.future import Future
@@ -361,7 +361,7 @@ def test_send_produce_request_raises_when_noleader(self, protocol, conn):
             "topic_noleader", 0,
             [create_message("a"), create_message("b")])]
 
-        with self.assertRaises(LeaderNotAvailableError):
+        with self.assertRaises(FailedPayloadsError):
             client.send_produce_request(requests)
 
     @patch('kafka.SimpleClient._get_conn')
@@ -386,7 +386,7 @@ def test_send_produce_request_raises_when_topic_unknown(self, protocol, conn):
             "topic_doesnt_exist", 0,
             [create_message("a"), create_message("b")])]
 
-        with self.assertRaises(UnknownTopicOrPartitionError):
+        with self.assertRaises(FailedPayloadsError):
             client.send_produce_request(requests)
 
     def test_timeout(self):
