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.

1 files changed, 18 insertions, 1 deletions

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)