check for broker None in send_broker_aware_request (added test for it)

Author: mrtheb

kafka/client.py

@@ -52,14 +52,22 @@ def _get_conn_for_broker(self, broker):
         return self.conns[(broker.host, broker.port)]
 
     def _get_leader_for_partition(self, topic, partition):
+        """
+        Returns the leader for a partition or None if the partition exists
+        but has no leader.
+
+        PartitionUnavailableError will be raised if the topic or partition
+        is not part of the metadata.
+        """
+
         key = TopicAndPartition(topic, partition)
         # reload metadata whether the partition is not available
-        # or has not leader (broker is None)
+        # or has no leader (broker is None)
         if self.topics_to_brokers.get(key) is None:
             self.load_metadata_for_topics(topic)
 
         if key not in self.topics_to_brokers:
-            raise PartitionUnavailableError("No leader for %s" % str(key))
+            raise PartitionUnavailableError("%s not available" % str(key))
 
         return self.topics_to_brokers[key]
 
@@ -115,8 +123,9 @@ def _send_broker_aware_request(self, payloads, encoder_fn, decoder_fn):
         for payload in payloads:
             leader = self._get_leader_for_partition(payload.topic,
                                                     payload.partition)
-            if leader == -1:
-                raise PartitionUnavailableError("Leader is unassigned for %s-%s" % payload.topic, payload.partition)
+            if leader is None:
+                raise PartitionUnavailableError(
+                    "No leader for topic %s partition %s" % (payload.topic, payload.partition))
             payloads_by_broker[leader].append(payload)
             original_keys.append((payload.topic, payload.partition))
 
@@ -249,7 +258,7 @@ def load_metadata_for_topics(self, *topics):
                 self.topic_partitions[topic].append(partition)
                 topic_part = TopicAndPartition(topic, partition)
                 if meta.leader == -1:
-                    log.info('No leader for topic %s partition %d', topic, partition)
+                    log.info('No leader for topic %s partition %s', topic, partition)
                     self.topics_to_brokers[topic_part] = None
                 else:
                     self.topics_to_brokers[topic_part] = brokers[meta.leader]

test/test_unit.py

@@ -517,5 +517,33 @@ def test_get_leader_returns_none_when_noleader(self, protocol, conn):
         self.assertEqual(brokers[0], client._get_leader_for_partition('topic_noleader', 0))
         self.assertEqual(brokers[1], client._get_leader_for_partition('topic_noleader', 1))
 
+    @patch('kafka.client.KafkaConnection')
+    @patch('kafka.client.KafkaProtocol')
+    def test_send_produce_request_raises_when_noleader(self, protocol, conn):
+        "Getting leader for partitions returns None when the partiion has no leader"
+
+        conn.recv.return_value = 'response'  # anything but None
+
+        brokers = {}
+        brokers[0] = BrokerMetadata(0, 'broker_1', 4567)
+        brokers[1] = BrokerMetadata(1, 'broker_2', 5678)
+
+        topics = {}
+        topics['topic_noleader'] = {
+            0: PartitionMetadata('topic_noleader', 0, -1, [], []),
+            1: PartitionMetadata('topic_noleader', 1, -1, [], [])
+        }
+        protocol.decode_metadata_response.return_value = (brokers, topics)
+
+        client = KafkaClient(host='broker_1', port=4567)
+
+        requests = [ProduceRequest(
+            "topic_noleader", 0,
+            [create_message("a"), create_message("b")])]
+
+        self.assertRaises(
+            PartitionUnavailableError,
+            client.send_produce_request, requests)
+
 if __name__ == '__main__':
     unittest.main()