Just encode to bytes if it isn't bytes.

This commit is contained in:
John Anderson
2015-01-24 16:32:09 -08:00
parent 3a52d18faf
commit 72bec18add
2 changed files with 14 additions and 8 deletions

View File

@@ -70,10 +70,12 @@ class SimpleProducer(Producer):
def send_messages(self, topic, *msg): def send_messages(self, topic, *msg):
if not isinstance(topic, six.binary_type): if not isinstance(topic, six.binary_type):
raise TypeError("topic must be type bytes") topic = topic.encode('utf-8')
partition = self._next_partition(topic) partition = self._next_partition(topic)
return super(SimpleProducer, self).send_messages(topic, partition, *msg) return super(SimpleProducer, self).send_messages(
topic, partition, *msg
)
def __repr__(self): def __repr__(self):
return '<SimpleProducer batch=%s>' % self.async return '<SimpleProducer batch=%s>' % self.async

View File

@@ -7,6 +7,7 @@ from . import unittest
from kafka.producer.base import Producer from kafka.producer.base import Producer
class TestKafkaProducer(unittest.TestCase): class TestKafkaProducer(unittest.TestCase):
def test_producer_message_types(self): def test_producer_message_types(self):
@@ -28,11 +29,14 @@ class TestKafkaProducer(unittest.TestCase):
def test_topic_message_types(self): def test_topic_message_types(self):
from kafka.producer.simple import SimpleProducer from kafka.producer.simple import SimpleProducer
producer = SimpleProducer(MagicMock()) client = MagicMock()
topic = "test-topic"
partition = 0
def send_message(): def partitions(topic):
producer.send_messages(topic, partition, b'hi') return [0, 1]
self.assertRaises(TypeError, send_message) client.get_partition_ids_for_topic = partitions
producer = SimpleProducer(client, random_start=False)
topic = b"test-topic"
producer.send_messages(topic, b'hi')
assert client.send_produce_request.called