Improve error handling and tests w.r.t. codecs

Add function kafka.protocol.create_message_set() that takes a list of
payloads and a codec and returns a message set with the desired encoding.

Introduce kafka.common.UnsupportedCodecError, raised if an unknown codec
is specified.

Include a test for the new function.

kafka/common.py

@@ -170,6 +170,11 @@ class ConsumerNoMoreData(KafkaError):
 class ProtocolError(KafkaError):
     pass
 
+
+class UnsupportedCodecError(KafkaError):
+    pass
+
+
 kafka_errors = {
     -1 : UnknownError,
     1  : OffsetOutOfRangeError,
@@ -187,6 +192,7 @@ kafka_errors = {
     13 : StaleLeaderEpochCodeError,
 }
 
+
 def check_error(response):
     error = kafka_errors.get(response.error)
     if error:

kafka/producer.py

@@ -9,12 +9,11 @@ from collections import defaultdict
 from itertools import cycle
 from multiprocessing import Queue, Process
 
-from kafka.common import ProduceRequest, TopicAndPartition
+from kafka.common import (
-from kafka.partitioner import HashedPartitioner
+    ProduceRequest, TopicAndPartition, UnsupportedCodecError
-from kafka.protocol import (
-    CODEC_NONE, CODEC_GZIP, CODEC_SNAPPY, ALL_CODECS,
-    create_message, create_gzip_message, create_snappy_message,
 )
+from kafka.partitioner import HashedPartitioner
+from kafka.protocol import CODEC_NONE, ALL_CODECS, create_message_set
 
 log = logging.getLogger("kafka")
 
@@ -66,13 +65,7 @@ def _send_upstream(queue, client, codec, batch_time, batch_size,
         # Send collected requests upstream
         reqs = []
         for topic_partition, msg in msgset.items():
-            if codec == CODEC_GZIP:
+            messages = create_message_set(msg, codec)
-                messages = [create_gzip_message(msg)]
-            elif codec == CODEC_SNAPPY:
-                messages = [create_snappy_message(msg)]
-            else:
-                messages = [create_message(m) for m in msg]
-
             req = ProduceRequest(topic_partition.topic,
                                  topic_partition.partition,
                                  messages)
@@ -132,7 +125,9 @@ class Producer(object):
 
         if codec is None:
             codec = CODEC_NONE
-        assert codec in ALL_CODECS
+        elif codec not in ALL_CODECS:
+            raise UnsupportedCodecError("Codec 0x%02x unsupported" % codec)
+
         self.codec = codec
 
         if self.async:
@@ -159,13 +154,7 @@ class Producer(object):
                 self.queue.put((TopicAndPartition(topic, partition), m))
             resp = []
         else:
-            if self.codec == CODEC_GZIP:
+            messages = create_message_set(msg, self.codec)
-                messages = [create_gzip_message(msg)]
-            elif self.codec == CODEC_SNAPPY:
-                messages = [create_snappy_message(msg)]
-            else:
-                messages = [create_message(m) for m in msg]
-
             req = ProduceRequest(topic, partition, messages)
             try:
                 resp = self.client.send_produce_request([req], acks=self.req_acks,

kafka/protocol.py

@@ -9,7 +9,8 @@ from kafka.common import (
     BrokerMetadata, PartitionMetadata, Message, OffsetAndMessage,
     ProduceResponse, FetchResponse, OffsetResponse,
     OffsetCommitResponse, OffsetFetchResponse, ProtocolError,
-    BufferUnderflowError, ChecksumError, ConsumerFetchSizeTooSmall
+    BufferUnderflowError, ChecksumError, ConsumerFetchSizeTooSmall,
+    UnsupportedCodecError
 )
 from kafka.util import (
     read_short_string, read_int_string, relative_unpack,
@@ -568,3 +569,19 @@ def create_snappy_message(payloads, key=None):
     codec = ATTRIBUTE_CODEC_MASK & CODEC_SNAPPY
 
     return Message(0, 0x00 | codec, key, snapped)
+
+
+def create_message_set(messages, codec=CODEC_NONE):
+    """Create a message set using the given codec.
+
+    If codec is CODEC_NONE, return a list of raw Kafka messages. Otherwise,
+    return a list containing a single codec-encoded message.
+    """
+    if codec == CODEC_NONE:
+        return [create_message(m) for m in messages]
+    elif codec == CODEC_GZIP:
+        return [create_gzip_message(messages)]
+    elif codec == CODEC_SNAPPY:
+        return [create_snappy_message(messages)]
+    else:
+        raise UnsupportedCodecError("Codec 0x%02x unsupported" % codec)

test/test_protocol.py

@@ -1,23 +1,30 @@
+import contextlib
+from contextlib import contextmanager
 import struct
 import unittest2
 
+import mock
+from mock import sentinel
+
 from kafka import KafkaClient
 from kafka.common import (
     OffsetRequest, OffsetCommitRequest, OffsetFetchRequest,
     OffsetResponse, OffsetCommitResponse, OffsetFetchResponse,
     ProduceRequest, FetchRequest, Message, ChecksumError,
-    ConsumerFetchSizeTooSmall, ProduceResponse, FetchResponse,
+    ConsumerFetchSizeTooSmall, ProduceResponse, FetchResponse, OffsetAndMessage,
-    OffsetAndMessage, BrokerMetadata, PartitionMetadata,
+    BrokerMetadata, PartitionMetadata, TopicAndPartition, KafkaUnavailableError,
-    TopicAndPartition, KafkaUnavailableError, ProtocolError,
+    ProtocolError, LeaderUnavailableError, PartitionUnavailableError,
-    LeaderUnavailableError, PartitionUnavailableError
+    UnsupportedCodecError
 )
 from kafka.codec import (
     has_snappy, gzip_encode, gzip_decode,
     snappy_encode, snappy_decode
 )
+import kafka.protocol
 from kafka.protocol import (
-    create_gzip_message, create_message, create_snappy_message, KafkaProtocol,
+    ATTRIBUTE_CODEC_MASK, CODEC_NONE, CODEC_GZIP, CODEC_SNAPPY, KafkaProtocol,
-    ATTRIBUTE_CODEC_MASK, CODEC_GZIP, CODEC_SNAPPY
+    create_message, create_gzip_message, create_snappy_message,
+    create_message_set
 )
 
 class TestProtocol(unittest2.TestCase):
@@ -691,3 +698,48 @@ class TestProtocol(unittest2.TestCase):
             OffsetFetchResponse(topic = 'topic1', partition = 2, offset = 4, error = 0, metadata = "meta"),
             OffsetFetchResponse(topic = 'topic1', partition = 4, offset = 8, error = 0, metadata = "meta"),
         ]))
+
+    @contextmanager
+    def mock_create_message_fns(self):
+        patches = contextlib.nested(
+            mock.patch.object(kafka.protocol, "create_message",
+                              return_value=sentinel.message),
+            mock.patch.object(kafka.protocol, "create_gzip_message",
+                              return_value=sentinel.gzip_message),
+            mock.patch.object(kafka.protocol, "create_snappy_message",
+                              return_value=sentinel.snappy_message),
+        )
+
+        with patches:
+            yield
+
+    def test_create_message_set(self):
+        messages = [1, 2, 3]
+
+        # Default codec is CODEC_NONE. Expect list of regular messages.
+        expect = [sentinel.message] * len(messages)
+        with self.mock_create_message_fns():
+            message_set = create_message_set(messages)
+        self.assertEqual(message_set, expect)
+
+        # CODEC_NONE: Expect list of regular messages.
+        expect = [sentinel.message] * len(messages)
+        with self.mock_create_message_fns():
+            message_set = create_message_set(messages, CODEC_NONE)
+        self.assertEqual(message_set, expect)
+
+        # CODEC_GZIP: Expect list of one gzip-encoded message.
+        expect = [sentinel.gzip_message]
+        with self.mock_create_message_fns():
+            message_set = create_message_set(messages, CODEC_GZIP)
+        self.assertEqual(message_set, expect)
+
+        # CODEC_SNAPPY: Expect list of one snappy-encoded message.
+        expect = [sentinel.snappy_message]
+        with self.mock_create_message_fns():
+            message_set = create_message_set(messages, CODEC_SNAPPY)
+        self.assertEqual(message_set, expect)
+
+        # Unknown codec should raise UnsupportedCodecError.
+        with self.assertRaises(UnsupportedCodecError):
+            create_message_set(messages, -1)