Add basic keyed producer tests

1 files changed, 22 insertions, 2 deletions

diff --git a/test/test_producer.py b/test/test_producer.py
index 27272f6..297b265 100644
--- a/test/test_producer.py
+++ b/test/test_producer.py
@@ -7,7 +7,7 @@ import time
 from mock import MagicMock, patch
 from . import unittest
 
-from kafka import KafkaClient, SimpleProducer
+from kafka import KafkaClient, SimpleProducer, KeyedProducer
 from kafka.common import (
     AsyncProducerQueueFull, FailedPayloadsError, NotLeaderForPartitionError,
     ProduceResponse, RetryOptions, TopicAndPartition
@@ -33,7 +33,8 @@ class TestKafkaProducer(unittest.TestCase):
         topic = b"test-topic"
         partition = 0
 
-        bad_data_types = (u'你怎么样?', 12, ['a', 'list'], ('a', 'tuple'), {'a': 'dict'})
+        bad_data_types = (u'你怎么样?', 12, ['a', 'list'],
+                          ('a', 'tuple'), {'a': 'dict'}, None,)
         for m in bad_data_types:
             with self.assertRaises(TypeError):
                 logging.debug("attempting to send message of type %s", type(m))
@@ -44,6 +45,25 @@ class TestKafkaProducer(unittest.TestCase):
             # This should not raise an exception
             producer.send_messages(topic, partition, m)
 
+    def test_keyedproducer_message_types(self):
+        client = MagicMock()
+        client.get_partition_ids_for_topic.return_value = [0, 1]
+        producer = KeyedProducer(client)
+        topic = b"test-topic"
+        key = 'testkey'
+
+        bad_data_types = (u'你怎么样?', 12, ['a', 'list'],
+                          ('a', 'tuple'), {'a': 'dict'},)
+        for m in bad_data_types:
+            with self.assertRaises(TypeError):
+                logging.debug("attempting to send message of type %s", type(m))
+                producer.send_messages(topic, key, m)
+
+        good_data_types = (b'a string!', None,)
+        for m in good_data_types:
+            # This should not raise an exception
+            producer.send_messages(topic, key, m)
+
     def test_topic_message_types(self):
         client = MagicMock()