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.

3 files changed, 33 insertions, 21 deletions

diff --git a/kafka/common.py b/kafka/common.py
index d515532..209754b 100644
--- a/kafka/common.py
+++ b/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:
diff --git a/kafka/producer.py b/kafka/producer.py
index 9ecb341..8e40be5 100644
--- a/kafka/producer.py
+++ b/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.partitioner import HashedPartitioner
-from kafka.protocol import (
-    CODEC_NONE, CODEC_GZIP, CODEC_SNAPPY, ALL_CODECS,
-    create_message, create_gzip_message, create_snappy_message,
+from kafka.common import (
+    ProduceRequest, TopicAndPartition, UnsupportedCodecError
 )
+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_gzip_message(msg)]
-            elif codec == CODEC_SNAPPY:
-                messages = [create_snappy_message(msg)]
-            else:
-                messages = [create_message(m) for m in msg]
-
+            messages = create_message_set(msg, codec)
             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_gzip_message(msg)]
-            elif self.codec == CODEC_SNAPPY:
-                messages = [create_snappy_message(msg)]
-            else:
-                messages = [create_message(m) for m in msg]
-
+            messages = create_message_set(msg, self.codec)
             req = ProduceRequest(topic, partition, messages)
             try:
                 resp = self.client.send_produce_request([req], acks=self.req_acks,
diff --git a/kafka/protocol.py b/kafka/protocol.py
index 730ae6b..58661c7 100644
--- a/kafka/protocol.py
+++ b/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)