Merge pull request #158 from wizzat/add_tests

Improve Tests, fix connection error timeout, other issues

1 files changed, 4 insertions, 5 deletions

diff --git a/kafka/protocol.py b/kafka/protocol.py
index 25be023..7ec7946 100644
--- a/kafka/protocol.py
+++ b/kafka/protocol.py
@@ -8,7 +8,7 @@ from kafka.codec import (
 from kafka.common import (
     BrokerMetadata, PartitionMetadata, Message, OffsetAndMessage,
     ProduceResponse, FetchResponse, OffsetResponse,
-    OffsetCommitResponse, OffsetFetchResponse,
+    OffsetCommitResponse, OffsetFetchResponse, ProtocolError,
     BufferUnderflowError, ChecksumError, ConsumerFetchSizeTooSmall
 )
 from kafka.util import (
@@ -50,7 +50,7 @@ class KafkaProtocol(object):
                            request_key,          # ApiKey
                            0,                    # ApiVersion
                            correlation_id,       # CorrelationId
-                           len(client_id),
+                           len(client_id),       # ClientId size
                            client_id)            # ClientId
 
     @classmethod
@@ -68,8 +68,7 @@ class KafkaProtocol(object):
         message_set = ""
         for message in messages:
             encoded_message = KafkaProtocol._encode_message(message)
-            message_set += struct.pack('>qi%ds' % len(encoded_message), 0,
-                                       len(encoded_message), encoded_message)
+            message_set += struct.pack('>qi%ds' % len(encoded_message), 0, len(encoded_message), encoded_message)
         return message_set
 
     @classmethod
@@ -96,7 +95,7 @@ class KafkaProtocol(object):
             crc = zlib.crc32(msg)
             msg = struct.pack('>i%ds' % len(msg), crc, msg)
         else:
-            raise Exception("Unexpected magic number: %d" % message.magic)
+            raise ProtocolError("Unexpected magic number: %d" % message.magic)
         return msg
 
     @classmethod