Refactor MessageSet and Message into LegacyRecordBatch to later support v2 message format

1 files changed, 4 insertions, 2 deletions

diff --git a/kafka/protocol/legacy.py b/kafka/protocol/legacy.py
index 37145b7..b8f84e7 100644
--- a/kafka/protocol/legacy.py
+++ b/kafka/protocol/legacy.py
@@ -19,6 +19,7 @@ from kafka.structs import ConsumerMetadataResponse
 from kafka.util import (
     crc32, read_short_string, relative_unpack,
     write_int_string, group_by_topic_and_partition)
+from kafka.protocol.message import MessageSet
 
 
 log = logging.getLogger(__name__)
@@ -144,7 +145,7 @@ class KafkaProtocol(object):
                           magic=msg.magic, attributes=msg.attributes
                     )
                     partition_msgs.append((0, m.encode()))
-                topic_msgs.append((partition, partition_msgs))
+                topic_msgs.append((partition, MessageSet.encode(partition_msgs, prepend_size=False)))
             topics.append((topic, topic_msgs))
 
 
@@ -215,7 +216,8 @@ class KafkaProtocol(object):
         ]
 
     @classmethod
-    def decode_message_set(cls, messages):
+    def decode_message_set(cls, raw_data):
+        messages = MessageSet.decode(raw_data, bytes_to_read=len(raw_data))
         for offset, _, message in messages:
             if isinstance(message, kafka.protocol.message.Message) and message.is_compressed():
                 inner_messages = message.decompress()