# -*- coding: utf-8 -*-
# Copyright 2015 Cyan, Inc.
# Copyright 2017, 2018, 2019, 2021 Ciena Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Test the KafkaCodec class.
"""

import struct
from contextlib import contextmanager
from unittest import SkipTest, TestCase, mock
from unittest.mock import sentinel

from .. import kafkacodec
from ..codec import gzip_decode, has_snappy, snappy_decode
from ..common import (
    BrokerMetadata, ChecksumError, ConsumerFetchSizeTooSmall,
    ConsumerMetadataResponse, FetchRequest, FetchResponse, InvalidMessageError,
    Message, OffsetAndMessage, OffsetCommitRequest, OffsetCommitResponse,
    OffsetFetchRequest, OffsetFetchResponse, OffsetRequest, OffsetResponse,
    PartitionMetadata, ProduceRequest, ProduceResponse, ProtocolError,
    TopicMetadata, UnsupportedCodecError, _HeartbeatRequest,
    _HeartbeatResponse, _JoinGroupProtocolMetadata, _JoinGroupRequest,
    _JoinGroupRequestProtocol, _JoinGroupResponse, _JoinGroupResponseMember,
    _LeaveGroupRequest, _LeaveGroupResponse, _SyncGroupMemberAssignment,
    _SyncGroupRequest, _SyncGroupRequestMember, _SyncGroupResponse,
)
from ..kafkacodec import (
    ATTRIBUTE_CODEC_MASK, CODEC_GZIP, CODEC_NONE, CODEC_SNAPPY, KafkaCodec,
    _ReprRequest, create_gzip_message, create_message, create_message_set,
    create_snappy_message,
)
from .testutil import make_send_requests


def create_encoded_metadata_response(broker_data, topic_data):
    encoded = struct.pack('>ii', 3, len(broker_data))
    for node_id, broker in broker_data.items():
        assert isinstance(broker.host, (str, bytes)), "{!r} must be string".format(broker.host)
        encoded += struct.pack('>ih', node_id, len(broker.host))
        encoded += broker.host.encode('ascii')
        encoded += struct.pack('>i', broker.port)

    encoded += struct.pack('>i', len(topic_data))
    for topic, topic_metadata in topic_data.items():
        assert isinstance(topic, (str, bytes)), "{!r} must be string".format(topic)
        _, topic_err, partitions = topic_metadata
        encoded += struct.pack('>hh',  topic_err, len(topic))
        encoded += topic.encode('ascii')
        encoded += struct.pack('>i', len(partitions))
        for partition, metadata in partitions.items():
            encoded += struct.pack('>hiii',
                                   metadata.partition_error_code,
                                   partition, metadata.leader,
                                   len(metadata.replicas))
            for replica in metadata.replicas:
                encoded += struct.pack('>i', replica)
            encoded += struct.pack('>i', len(metadata.isr))
            for isr in metadata.isr:
                encoded += struct.pack('>i', isr)

    return encoded


class ReprRequestTests(TestCase):
    def test_produce(self):
        rr = _ReprRequest((
            b'\x00\x00'
            b'\x00\x03'
            b'\x00\x00\x04\x00'
            b'xx'
        ))
        self.assertEqual('ProduceRequest3 correlationId=1024 (10 bytes)', str(rr))

    def test_unknown_key(self):
        rr = _ReprRequest((
            b'\x00\xff'
            b'\x00\x01'
            b'\x00\x00\x00\x00'
        ))
        self.assertEqual('request key=255v1 correlationId=0 (8 bytes)', str(rr))

    def test_too_short(self):
        rr = _ReprRequest(b'\x00\x00\xab\xcd')
        self.assertEqual("invalid request (0000abcd)", str(rr))


class TestKafkaCodec(TestCase):
    def test_create_message(self):
        payload = b"test"
        key = b"key"
        msg = create_message(payload, key)
        self.assertEqual(msg.magic, 0)
        self.assertEqual(msg.attributes, 0)
        self.assertEqual(msg.key, key)
        self.assertEqual(msg.value, payload)

    def test_create_gzip(self):
        message_list = [create_message(b"v1", None),
                        create_message(b"v2", key=b'42')]
        msg = create_gzip_message(message_list)
        self.assertEqual(msg.magic, 0)
        self.assertEqual(msg.attributes, ATTRIBUTE_CODEC_MASK & CODEC_GZIP)
        self.assertEqual(msg.key, None)
        # Need to decode to check since gzipped payload is non-deterministic
        decoded = gzip_decode(msg.value)
        expect = b"".join([
            struct.pack(">q", 0),           # MsgSet offset
            struct.pack(">i", 16),          # MsgSet size
            struct.pack(">i", 1285512130),  # CRC
            struct.pack(">bb", 0, 0),       # Magic, flags
            struct.pack(">i", -1),          # -1 indicates a null key
            struct.pack(">i", 2),           # Msg length (bytes)
            b"v1",                          # Message contents

            struct.pack(">q", 0),           # MsgSet offset
            struct.pack(">i", 18),          # MsgSet size
            struct.pack(">i", 1929437987),  # CRC
            struct.pack(">bb", 0, 0),       # Magic, flags
            struct.pack(">i2s", 2, b'42'),  # Key is 2 bytes long, 42
            struct.pack(">i", 2),           # Msg length (bytes)
            b"v2",                          # Message contents
        ])

        self.assertEqual(decoded, expect)

    def test_create_snappy(self):
        if not has_snappy():
            raise SkipTest("Snappy not available")  # pragma: no cover
        message_list = [create_message(b"v3", key=b'84'),
                        create_message(b"v4", None)]
        msg = create_snappy_message(message_list)
        self.assertEqual(msg.magic, 0)
        self.assertEqual(msg.attributes, ATTRIBUTE_CODEC_MASK & CODEC_SNAPPY)
        self.assertEqual(msg.key, None)
        decoded = snappy_decode(msg.value)
        expect = b"".join([
            struct.pack(">q", 0),           # MsgSet offset
            struct.pack(">i", 18),          # MsgSet size
            struct.pack(">i", 813233088),   # CRC
            struct.pack(">bb", 0, 0),       # Magic, flags
            struct.pack(">i2s", 2, b'84'),  # Key is 2 bytes long, '84'
            struct.pack(">i", 2),           # Msg length (bytes)
            b"v3",                          # Message contents

            struct.pack(">q", 0),           # MsgSet offset
            struct.pack(">i", 16),          # MsgSet size
            struct.pack(">i", 1022734157),  # CRC
            struct.pack(">bb", 0, 0),       # Magic, flags
            struct.pack(">i", -1),          # -1 indicates a null key
            struct.pack(">i", 2),           # Msg length (bytes)
            b"v4",                          # Message contents
        ])

        self.assertEqual(decoded, expect)

    def test_encode_message_header(self):
        expect = b"".join([
            struct.pack(">h", 10),              # API Key
            struct.pack(">h", 0),               # API Version
            struct.pack(">i", 4),               # Correlation Id
            struct.pack(">h", len("client3")),  # Length of clientId
            b"client3",                         # ClientId
        ])

        encoded = KafkaCodec._encode_message_header(b"client3", 4, 10)
        self.assertEqual(encoded, expect)

    def test_encode_message(self):
        message = create_message(b"test", b"key")
        encoded = KafkaCodec._encode_message(message)
        expect = b"".join([
            struct.pack(">i", -1427009701),  # CRC
            struct.pack(">bb", 0, 0),        # Magic, flags
            struct.pack(">i", 3),            # Length of key
            b"key",                          # key
            struct.pack(">i", 4),            # Length of value
            b"test",                         # value
        ])

        self.assertEqual(encoded, expect)

    def test_decode_message(self):
        encoded = b"".join([
            struct.pack(">i", -1427009701),  # CRC
            struct.pack(">bb", 0, 0),        # Magic, flags
            struct.pack(">i", 3),            # Length of key
            b"key",                          # key
            struct.pack(">i", 4),            # Length of value
            b"test",                         # value
        ])

        offset = 10
        (returned_offset, decoded_message) = list(
            KafkaCodec._decode_message(encoded, offset))[0]

        self.assertEqual(returned_offset, offset)
        self.assertEqual(decoded_message, create_message(b"test", b"key"))

    def test_encode_message_failure(self):
        self.assertRaises(ProtocolError,
                          KafkaCodec._encode_message,
                          Message(1, 0, "key", "test"))

    def test_encode_message_set(self):
        message_set = [
            create_message(b"v1", b"k1"),
            create_message(b"v2", b"k2"),
        ]

        encoded = KafkaCodec._encode_message_set(message_set)
        expect = b"".join([
            struct.pack(">q", 0),           # MsgSet Offset
            struct.pack(">i", 18),          # Msg Size
            struct.pack(">i", 1474775406),  # CRC
            struct.pack(">bb", 0, 0),       # Magic, flags
            struct.pack(">i", 2),           # Length of key
            b"k1",                          # Key
            struct.pack(">i", 2),           # Length of value
            b"v1",                          # Value

            struct.pack(">q", 0),          # MsgSet Offset
            struct.pack(">i", 18),         # Msg Size
            struct.pack(">i", -16383415),  # CRC
            struct.pack(">bb", 0, 0),      # Magic, flags
            struct.pack(">i", 2),          # Length of key
            b"k2",                         # Key
            struct.pack(">i", 2),          # Length of value
            b"v2",                         # Value
        ])

        self.assertEqual(encoded, expect)

    def test_decode_message_set(self):
        encoded = b"".join([
            struct.pack(">q", 0),           # MsgSet Offset
            struct.pack(">i", 18),          # Msg Size
            struct.pack(">i", 1474775406),  # CRC
            struct.pack(">bb", 0, 0),       # Magic, flags
            struct.pack(">i", 2),           # Length of key
            b"k1",                          # Key
            struct.pack(">i", 2),           # Length of value
            b"v1",                          # Value

            struct.pack(">q", 1),          # MsgSet Offset
            struct.pack(">i", 18),         # Msg Size
            struct.pack(">i", -16383415),  # CRC
            struct.pack(">bb", 0, 0),      # Magic, flags
            struct.pack(">i", 2),          # Length of key
            b"k2",                         # Key
            struct.pack(">i", 2),          # Length of value
            b"v2",                         # Value
        ])

        msgs = list(KafkaCodec._decode_message_set_iter(encoded))
        self.assertEqual(len(msgs), 2)
        msg1, msg2 = msgs

        returned_offset1, decoded_message1 = msg1
        returned_offset2, decoded_message2 = msg2

        self.assertEqual(returned_offset1, 0)
        self.assertEqual(decoded_message1, create_message(b"v1", b"k1"))

        self.assertEqual(returned_offset2, 1)
        self.assertEqual(decoded_message2, create_message(b"v2", b"k2"))

    def test_decode_message_gzip(self):
        gzip_encoded = (b'\xc0\x11\xb2\xf0\x00\x01\xff\xff\xff\xff\x00\x00\x000'
                        b'\x1f\x8b\x08\x00\xa1\xc1\xc5R\x02\xffc`\x80\x03\x01'
                        b'\x9f\xf9\xd1\x87\x18\x18\xfe\x03\x01\x90\xc7Tf\xc8'
                        b'\x80$wu\x1aW\x05\x92\x9c\x11\x00z\xc0h\x888\x00\x00'
                        b'\x00')
        offset = 11
        messages = list(KafkaCodec._decode_message(gzip_encoded, offset))

        self.assertEqual(len(messages), 2)
        msg1, msg2 = messages

        returned_offset1, decoded_message1 = msg1
        self.assertEqual(returned_offset1, 0)
        self.assertEqual(decoded_message1, create_message(b"v1"))

        returned_offset2, decoded_message2 = msg2
        self.assertEqual(returned_offset2, 0)
        self.assertEqual(decoded_message2, create_message(b"v2"))

    def test_decode_message_snappy(self):
        if not has_snappy():
            raise SkipTest("Snappy not available")  # pragma: no cover
        snappy_encoded = (b'\xec\x80\xa1\x95\x00\x02\xff\xff\xff\xff\x00\x00'
                          b'\x00,8\x00\x00\x19\x01@\x10L\x9f[\xc2\x00\x00\xff'
                          b'\xff\xff\xff\x00\x00\x00\x02v1\x19\x1bD\x00\x10\xd5'
                          b'\x96\nx\x00\x00\xff\xff\xff\xff\x00\x00\x00\x02v2')
        offset = 11
        messages = list(KafkaCodec._decode_message(snappy_encoded, offset))
        self.assertEqual(len(messages), 2)

        msg1, msg2 = messages

        returned_offset1, decoded_message1 = msg1
        self.assertEqual(returned_offset1, 0)
        self.assertEqual(decoded_message1, create_message(b"v1"))

        returned_offset2, decoded_message2 = msg2
        self.assertEqual(returned_offset2, 0)
        self.assertEqual(decoded_message2, create_message(b"v2"))

    def test_decode_message_checksum_error(self):
        invalid_encoded_message = b"This is not a valid encoded message"
        iter = KafkaCodec._decode_message(invalid_encoded_message, 0)
        self.assertRaises(ChecksumError, list, iter)

    # NOTE: The error handling in _decode_message_set_iter() is questionable.
    # If it's modified, the next two tests might need to be fixed.
    def test_decode_message_set_fetch_size_too_small(self):
        self.assertRaises(ConsumerFetchSizeTooSmall,
                          list, KafkaCodec._decode_message_set_iter('a'))

    def test_decode_message_set_stop_iteration(self):
        encoded = b"".join([
            struct.pack(">q", 0),           # MsgSet Offset
            struct.pack(">i", 18),          # Msg Size
            struct.pack(">i", 1474775406),  # CRC
            struct.pack(">bb", 0, 0),       # Magic, flags
            struct.pack(">i", 2),           # Length of key
            b"k1",                          # Key
            struct.pack(">i", 2),           # Length of value
            b"v1",                          # Value

            struct.pack(">q", 1),           # MsgSet Offset
            struct.pack(">i", 18),          # Msg Size
            struct.pack(">i", -16383415),   # CRC
            struct.pack(">bb", 0, 0),       # Magic, flags
            struct.pack(">i", 2),           # Length of key
            b"k2",                          # Key
            struct.pack(">i", 2),           # Length of value
            b"v2",                          # Value
            b"@1$%(Y!",                     # Random padding
        ])

        msgs = list(KafkaCodec._decode_message_set_iter(encoded))
        self.assertEqual(len(msgs), 2)
        msg1, msg2 = msgs

        returned_offset1, decoded_message1 = msg1
        returned_offset2, decoded_message2 = msg2

        self.assertEqual(returned_offset1, 0)
        self.assertEqual(decoded_message1, create_message(b"v1", b"k1"))

        self.assertEqual(returned_offset2, 1)
        self.assertEqual(decoded_message2, create_message(b"v2", b"k2"))

    def test_get_response_correlation_id(self):
        t1 = b"topic1"
        t2 = b"topic2"
        corrID = 30372
        encoded = struct.pack('>iih%dsiihqihqh%dsiihq' % (len(t1), len(t2)),
                              corrID, 2, len(t1), t1, 2, 0, 0, 10, 1, 1, 20,
                              len(t2), t2, 1, 0, 0, 30)
        self.assertEqual(
            corrID, KafkaCodec.get_response_correlation_id(encoded))

    def test_encode_produce_request(self):
        requests = [
            ProduceRequest("topic1", 0, [
                create_message(b"a"),
                create_message(b"b"),
            ]),
            ProduceRequest(u"topic2", 1, [
                create_message(b"c"),
            ]),
        ]

        msg_a_binary = KafkaCodec._encode_message(create_message(b"a"))
        msg_b_binary = KafkaCodec._encode_message(create_message(b"b"))
        msg_c_binary = KafkaCodec._encode_message(create_message(b"c"))

        header = b"".join([
            struct.pack('>h', 0),        # Msg Header, Message type = Produce
            struct.pack('>h', 0),        # Msg Header, API version
            struct.pack('>i', 2),        # Msg Header, Correlation ID
            struct.pack('>h7s', 7, b"client1"),  # Msg Header, The client ID
            struct.pack('>h', 2),        # Num acks required
            struct.pack('>i', 100),      # Request Timeout
            struct.pack('>i', 2),        # The number of requests
        ])

        total_len = len(msg_a_binary) + len(msg_b_binary)
        topic1 = b"".join([
            struct.pack('>h6s', 6, b'topic1'),         # The topic1
            struct.pack('>i', 1),                      # One message set
            struct.pack('>i', 0),                      # Partition 0
            struct.pack('>i', total_len + 24),         # Size of message set
            struct.pack('>q', 0),                      # No offset specified
            struct.pack('>i', len(msg_a_binary)),      # Length of message
            msg_a_binary,                              # Actual message
            struct.pack('>q', 0),                      # No offset specified
            struct.pack('>i', len(msg_b_binary)),      # Length of message
            msg_b_binary,                              # Actual message
        ])

        topic2 = b"".join([
            struct.pack('>h6s', 6, b'topic2'),          # The topic1
            struct.pack('>i', 1),                       # One message set
            struct.pack('>i', 1),                       # Partition 1
            struct.pack('>i', len(msg_c_binary) + 12),  # Size of message set
            struct.pack('>q', 0),                       # No offset specified
            struct.pack('>i', len(msg_c_binary)),       # Length of message
            msg_c_binary,                               # Actual message
        ])

        expected1 = b"".join([header, topic1, topic2])
        expected2 = b"".join([header, topic2, topic1])

        encoded = KafkaCodec.encode_produce_request(
            b"client1", 2, requests, 2, 100)
        self.assertIn(encoded, [expected1, expected2])

    def test_decode_produce_response(self):
        t1 = "topic1"
        t2 = u"topic2"
        encoded = struct.pack('>iih%dsiihqihqh%dsiihq' % (len(t1), len(t2)),
                              2, 2, len(t1), t1.encode(), 2, 0, 0, 10, 1, 1, 20,
                              len(t2), t2.encode(), 1, 0, 0, 30)
        responses = list(KafkaCodec.decode_produce_response(encoded))
        self.assertEqual(responses,
                         [ProduceResponse(t1, 0, 0, 10),
                          ProduceResponse(t1, 1, 1, 20),
                          ProduceResponse(t2, 0, 0, 30)])

    def test_encode_fetch_request(self):
        requests = [
            FetchRequest("topic1", 0, 10, 1024),
            FetchRequest(u"topic2", 1, 20, 100),
        ]

        header = b"".join([
            struct.pack('>h', 1),       # Msg Header, Message type = Fetch
            struct.pack('>h', 0),               # Msg Header, API version
            struct.pack('>i', 3),               # Msg Header, Correlation ID
            struct.pack('>h7s', 7, b"client1"),  # Msg Header, The client ID
            struct.pack('>i', -1),              # Replica Id
            struct.pack('>i', 2),               # Max wait time
            struct.pack('>i', 100),             # Min bytes
            struct.pack('>i', 2),               # Num requests
        ])

        topic1 = b"".join([
            struct.pack('>h6s', 6, b'topic1'),  # Topic
            struct.pack('>i', 1),              # Num Payloads
            struct.pack('>i', 0),              # Partition 0
            struct.pack('>q', 10),             # Offset
            struct.pack('>i', 1024),           # Max Bytes
        ])

        topic2 = b"".join([
            struct.pack('>h6s', 6, b'topic2'),  # Topic
            struct.pack('>i', 1),              # Num Payloads
            struct.pack('>i', 1),              # Partition 0
            struct.pack('>q', 20),             # Offset
            struct.pack('>i', 100),            # Max Bytes
        ])

        expected1 = b"".join([header, topic1, topic2])
        expected2 = b"".join([header, topic2, topic1])

        encoded = KafkaCodec.encode_fetch_request(
            b"client1", 3, requests, 2, 100)
        self.assertIn(encoded, [expected1, expected2])

    def test_decode_fetch_response(self):
        t1 = "topic1"
        t2 = u"topic2"
        msgs = [create_message(m) for m in [b"message1", b"hi", b"boo", b"foo", b"so fun!"]]
        ms1 = KafkaCodec._encode_message_set([msgs[0], msgs[1]])
        ms2 = KafkaCodec._encode_message_set([msgs[2]])
        ms3 = KafkaCodec._encode_message_set([msgs[3], msgs[4]])

        encoded = struct.pack('>iih%dsiihqi%dsihqi%dsh%dsiihqi%ds' %
                              (len(t1), len(ms1), len(ms2), len(t2), len(ms3)),
                              4, 2, len(t1), t1.encode(), 2, 0, 0, 10, len(ms1), ms1, 1,
                              1, 20, len(ms2), ms2, len(t2), t2.encode(), 1, 0, 0, 30,
                              len(ms3), ms3)

        responses = list(KafkaCodec.decode_fetch_response(encoded))

        def expand_messages(response):
            return FetchResponse(response.topic, response.partition,
                                 response.error, response.highwaterMark,
                                 list(response.messages))

        expanded_responses = [expand_messages(r) for r in responses]
        expect = [FetchResponse(t1, 0, 0, 10, [OffsetAndMessage(0, msgs[0]),
                                               OffsetAndMessage(0, msgs[1])]),
                  FetchResponse(t1, 1, 1, 20, [OffsetAndMessage(0, msgs[2])]),
                  FetchResponse(t2, 0, 0, 30, [OffsetAndMessage(0, msgs[3]),
                                               OffsetAndMessage(0, msgs[4])])]
        self.assertEqual(expanded_responses, expect)

    def test_encode_metadata_request_no_topics(self):
        expected = b"".join([
            struct.pack('>h', 3),           # API key metadata fetch
            struct.pack('>h', 0),           # API version
            struct.pack('>i', 4),           # Correlation ID
            struct.pack('>h3s', 3, b"cid"),  # The client ID
            struct.pack('>i', 0),           # No topics, give all the data!
        ])

        encoded = KafkaCodec.encode_metadata_request(b"cid", 4)

        self.assertEqual(encoded, expected)

    def test_encode_metadata_request_with_topics(self):
        expected = b"".join([
            struct.pack('>h', 3),           # API key metadata fetch
            struct.pack('>h', 0),           # API version
            struct.pack('>i', 4),           # Correlation ID
            struct.pack('>h3s', 3, b"cid"),  # The client ID
            struct.pack('>i', 2),           # Number of topics in the request
            struct.pack('>h2s', 2, b"t1"),   # Topic "t1"
            struct.pack('>h2s', 2, b"t2"),   # Topic "t2"
        ])

        encoded = KafkaCodec.encode_metadata_request(b"cid", 4, [u"t1", "t2"])

        self.assertEqual(encoded, expected)

    def test_decode_metadata_response(self):
        node_brokers = {
            0: BrokerMetadata(0, "brokers1.afkak.rdio.com", 1000),
            1: BrokerMetadata(1, "brokers1.afkak.rdio.com", 1001),
            3: BrokerMetadata(3, "brokers2.afkak.rdio.com", 1000),
        }

        topic_partitions = {
            "topic1": TopicMetadata(
                'topic1', 0, {
                    0: PartitionMetadata(u"topic1", 0, 0, 1, (0, 2), (2,)),
                    1: PartitionMetadata("topic1", 1, 1, 3, (0, 1), (0, 1)),
                },
            ),
            u"topic2": TopicMetadata(
                u'topic2', 1, {
                    0: PartitionMetadata(u"topic2", 0, 0, 0, (), ()),
                },
            ),
        }
        encoded = create_encoded_metadata_response(
            node_brokers, topic_partitions)
        decoded = KafkaCodec.decode_metadata_response(encoded)
        self.assertEqual(decoded, (node_brokers, topic_partitions))

    def test_decode_metadata_response_failures(self):
        data = b"".join([
            struct.pack(">ii", 0xCAFE, 0xBABE),  # correlation ID & numbrokers
            struct.pack('>i', 1234),             # Node Id
            struct.pack('>h', len(b"hostname")), b"hostname",  # host
            struct.pack('>i', 1025),             # port
        ])
        self.assertRaises(InvalidMessageError,
                          KafkaCodec.decode_metadata_response, data)

    def test_encode_offset_request(self):
        expected = b"".join([
            struct.pack('>h', 2),            # Message type = offset fetch
            struct.pack('>h', 0),            # API version
            struct.pack('>i', 4),            # Correlation ID
            struct.pack('>h3s', 3, b"cid"),  # The client ID
            struct.pack('>i', -1),           # Replica Id
            struct.pack('>i', 0),            # No topic/partitions
        ])

        encoded = KafkaCodec.encode_offset_request(b"cid", 4)

        self.assertEqual(encoded, expected)

    def test_encode_offset_request__no_payload(self):
        expected = b"".join([
            struct.pack('>h', 2),             # Message type = offset fetch
            struct.pack('>h', 0),             # API version
            struct.pack('>i', 4),             # Correlation ID
            struct.pack('>h3s', 3, b"cid"),   # The client ID
            struct.pack('>i', -1),            # Replica Id
            struct.pack('>i', 1),             # Num topics
            struct.pack(">h6s", 6, b"topic1"),  # Topic for the request
            struct.pack(">i", 2),             # Two partitions

            struct.pack(">i", 3),             # Partition 3
            struct.pack(">q", -1),            # No time offset
            struct.pack(">i", 1),             # One offset requested

            struct.pack(">i", 4),             # Partition 3
            struct.pack(">q", -1),            # No time offset
            struct.pack(">i", 1),             # One offset requested
        ])

        encoded = KafkaCodec.encode_offset_request(b"cid", 4, [
            OffsetRequest('topic1', 3, -1, 1),
            OffsetRequest(u'topic1', 4, -1, 1),
        ])

        self.assertEqual(encoded, expected)

    def test_decode_offset_response(self):
        encoded = b"".join([
            struct.pack(">i", 42),            # Correlation ID
            struct.pack(">i", 1),             # One topics
            struct.pack(">h6s", 6, b"topic1"),  # First topic
            struct.pack(">i", 2),             # Two partitions

            struct.pack(">i", 2),             # Partition 2
            struct.pack(">h", 0),             # No error
            struct.pack(">i", 1),             # One offset
            struct.pack(">q", 4),             # Offset 4

            struct.pack(">i", 4),             # Partition 4
            struct.pack(">h", 0),             # No error
            struct.pack(">i", 1),             # One offset
            struct.pack(">q", 8),             # Offset 8
        ])

        results = KafkaCodec.decode_offset_response(encoded)
        self.assertEqual(set(results), set([
            OffsetResponse(topic='topic1', partition=2, error=0, offsets=(4,)),
            OffsetResponse(topic=u'topic1', partition=4, error=0, offsets=(8,)),
        ]))

    def test_encode_offset_commit_request(self):
        header = b"".join([
            struct.pack('>h', 8),                # Message type = offset commit
            struct.pack('>h', 1),                # API version
            struct.pack('>i', 42),               # Correlation ID
            struct.pack('>h9s', 9, b"client_id"),  # The client ID
            struct.pack('>h8s', 8, b"group_id"),  # The group to commit for
            struct.pack('>i', 996),              # Group generation ID
            struct.pack('>h11s', 11, b'consumer_id'),  # Consumer ID
            struct.pack('>i', 2),                # Num topics
        ])

        topic1 = b"".join([
            struct.pack(">h6s", 6, b"topic1"),    # Topic for the request
            struct.pack(">i", 2),                # Two partitions
            struct.pack(">i", 0),                # Partition 0
            struct.pack(">q", 123),              # Offset 123
            struct.pack(">q", 1437585816816),    # Timestamp in ms > epoch
            struct.pack(">h", -1),               # Null metadata
            struct.pack(">i", 1),                # Partition 1
            struct.pack(">q", 234),              # Offset 234
            struct.pack(">q", 1436981054199),    # Timestamp in ms > epoch
            struct.pack(">h11s", 11, b'My_Metadata'),  # Null metadata
        ])

        topic2 = b"".join([
            struct.pack(">h6s", 6, b"topic2"),    # Topic for the request
            struct.pack(">i", 1),                # One partition
            struct.pack(">i", 2),                # Partition 2
            struct.pack(">q", 345),              # Offset 345
            struct.pack(">q", -1),               # Timestamp 'invalid-time'
            struct.pack(">h", -1),               # Null metadata
        ])

        # A dict is used, so we can't predict the order of the topics...
        expected1 = b"".join([header, topic1, topic2])
        expected2 = b"".join([header, topic2, topic1])

        encoded = KafkaCodec.encode_offset_commit_request(
            b"client_id", 42, u"group_id", 996, u'consumer_id', [
                OffsetCommitRequest("topic1", 0, 123, 1437585816816, None),
                OffsetCommitRequest(u"topic1", 1, 234, 1436981054199,
                                    b'My_Metadata'),
                OffsetCommitRequest(u"topic2", 2, 345, -1, None),
            ])

        self.assertIn(encoded, [expected1, expected2])

    def test_decode_offset_commit_response(self):
        encoded = b"".join([
            struct.pack(">i", 42),            # Correlation ID
            struct.pack(">i", 1),             # One topic
            struct.pack(">h6s", 6, b"topic1"),  # First topic
            struct.pack(">i", 2),             # Two partitions

            struct.pack(">i", 2),             # Partition 2
            struct.pack(">h", 0),             # No error

            struct.pack(">i", 4),             # Partition 4
            struct.pack(">h", 0),             # No error
        ])

        results = KafkaCodec.decode_offset_commit_response(encoded)
        self.assertEqual(set(results), set([
            OffsetCommitResponse(topic='topic1', partition=2, error=0),
            OffsetCommitResponse(topic=u'topic1', partition=4, error=0),
        ]))

    def test_encode_offset_fetch_request(self):
        header = b"".join([
            struct.pack('>h', 9),                  # Message type = offset fetch
            struct.pack('>h', 1),                  # API version
            struct.pack('>i', 42),                 # Correlation ID
            struct.pack('>h9s', 9, b"client_id"),  # The client ID
            struct.pack('>h8s', 8, b"group_id"),   # The group to commit for
            struct.pack('>i', 2),                  # Num topics
        ])

        topic1 = b"".join([
            struct.pack(">h6s", 6, b"topic1"),   # Topic for the request
            struct.pack(">i", 2),                # Two partitions
            struct.pack(">i", 0),                # Partition 0
            struct.pack(">i", 1),                # Partition 1
        ])

        topic2 = b"".join([
            struct.pack(">h6s", 6, b"topic2"),   # Topic for the request
            struct.pack(">i", 1),                # One partitions
            struct.pack(">i", 2),                # Partition 2
        ])

        expected1 = b"".join([header, topic1, topic2])
        expected2 = b"".join([header, topic2, topic1])

        encoded = KafkaCodec.encode_offset_fetch_request(
            b"client_id", 42, u"group_id", [
                OffsetFetchRequest("topic1", 0),
                OffsetFetchRequest(u"topic1", 1),
                OffsetFetchRequest("topic2", 2),
            ])

        self.assertIn(encoded, [expected1, expected2])

    def test_decode_offset_fetch_response(self):
        encoded = b"".join([
            struct.pack(">i", 42),              # Correlation ID
            struct.pack(">i", 1),               # One topics
            struct.pack(">h6s", 6, b"topic1"),  # First topic
            struct.pack(">i", 2),               # Two partitions

            struct.pack(">i", 2),             # Partition 2
            struct.pack(">q", 4),             # Offset 4
            struct.pack(">h4s", 4, b"meta"),  # Metadata
            struct.pack(">h", 0),             # No error

            struct.pack(">i", 4),             # Partition 4
            struct.pack(">q", 8),             # Offset 8
            struct.pack(">h4s", 4, b"meta"),  # Metadata
            struct.pack(">h", 0),             # No error
        ])

        results = KafkaCodec.decode_offset_fetch_response(encoded)
        self.assertEqual(set(results), set([
            OffsetFetchResponse(topic=u'topic1', partition=2, offset=4,
                                error=0, metadata=b"meta"),
            OffsetFetchResponse(topic='topic1', partition=4, offset=8,
                                error=0, metadata=b"meta"),
        ]))

    @contextmanager
    def mock_create_message_fns(self):
        p1 = mock.patch.object(kafkacodec, "create_message",
                               return_value=sentinel.message)
        p2 = mock.patch.object(kafkacodec, "create_gzip_message",
                               return_value=sentinel.gzip_message)
        p3 = mock.patch.object(kafkacodec, "create_snappy_message",
                               return_value=sentinel.snappy_message)
        with p1, p2, p3:
            yield

    def test_create_message_set(self):
        messages = [b'1', b'2', b'3']
        reqs = make_send_requests(messages)

        # 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(reqs)
        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(reqs, 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(reqs, 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(reqs, CODEC_SNAPPY)
        self.assertEqual(message_set, expect)

        # Unknown codec should raise UnsupportedCodecError.
        self.assertRaises(UnsupportedCodecError,
                          create_message_set, reqs, -1)

    def test_encode_consumer_metadata_request(self):
        expected = b"".join([
            struct.pack('>h', 10),          # API key ConsumerMetadataRequest
            struct.pack('>h', 0),           # API version
            struct.pack('>i', 9),           # Correlation ID
            struct.pack('>h3s', 3, b"cID"),  # The client ID
            struct.pack('>h6s', 6, b"group1"),   # Consumer group 'group1'
        ])

        encoded = KafkaCodec.encode_consumermetadata_request(b"cID", 9, u"group1")

        self.assertEqual(encoded, expected)

    def test_decode_consumermetadata_response(self):
        expected = ConsumerMetadataResponse(0, 5, "theHost", 4242)
        encoded = b"".join([
            struct.pack('>i', 9),  # Correlation ID
            struct.pack('>h', 0),  # Error Code
            struct.pack('>i', 5),  # Coordinator ID
            struct.pack('>h', len(b"theHost")), b"theHost",  # Coordinator Host
            struct.pack('>i', 4242),  # Coordinator port
        ])
        decoded = KafkaCodec.decode_consumermetadata_response(encoded)
        self.assertEqual(decoded, expected)

    def test_encode_join_group_protocol_metadata(self):
        expected = b"".join([
            struct.pack('>h', 1),                    # version 1
            struct.pack('>ih5s', 1, 5, b"topic"),    # one topic
            struct.pack('>i4s', 4, b"data"),         # Metadata
        ])

        encoded = KafkaCodec.encode_join_group_protocol_metadata(
            1,
            ["topic"],
            b"data",
        )
        self.assertEqual(encoded, expected)

    def test_decode_join_group_protocol_metadata(self):
        expected = _JoinGroupProtocolMetadata(1, ["topic"], b"data")

        encoded = b"".join([
            struct.pack('>h', 1),                    # version 1
            struct.pack('>ih5s', 1, 5, b"topic"),    # one topic
            struct.pack('>i4s', 4, b"data"),         # Metadata
        ])
        decoded = KafkaCodec.decode_join_group_protocol_metadata(encoded)
        self.assertEqual(decoded, expected)

    def test_encode_join_group_request(self):
        expected = b"".join([
            struct.pack('>h', 11),          # API key JoinGroupRequest
            struct.pack('>h', 0),           # API version
            struct.pack('>i', 1),           # Correlation ID
            struct.pack('>h3s', 3, b"cID"),      # The client ID
            struct.pack('>h6s', 6, b"group1"),   # Consumer group 'group1'
            struct.pack('>i', 1),                # Timeout
            struct.pack('>h6s', 6, b"member"),   # Member ID
            struct.pack('>h5s', 5, b"proto"),    # Protocol Type
            struct.pack('>i', 2),                # Two protocols

            struct.pack('>h4s', 4, b"name"),     # Name
            struct.pack('>i4s', 4, b"meta"),     # Metadata
            struct.pack('>h5s', 5, b"name2"),    # Name
            struct.pack('>i5s', 5, b"meta2"),    # Metadata
        ])

        encoded = KafkaCodec.encode_join_group_request(
            b"cID", 1,
            _JoinGroupRequest("group1", 1, "member", "proto", [
                _JoinGroupRequestProtocol("name", b"meta"),
                _JoinGroupRequestProtocol("name2", b"meta2"),
            ]),
        )
        self.assertEqual(encoded, expected)

    def test_decode_join_group_response(self):
        expected = _JoinGroupResponse(
            0, 1, "proto",
            "leader", "member", [
                _JoinGroupResponseMember("id1", b"data1"),
                _JoinGroupResponseMember("id2", b"data2"),
            ])

        encoded = b"".join([
            struct.pack('>i', 9),  # Correlation ID
            struct.pack('>h', 0),  # Error Code
            struct.pack('>i', 1),  # Generation ID
            struct.pack('>h5s', 5, b"proto"),   # Group Protocol
            struct.pack('>h6s', 6, b"leader"),  # Leader ID
            struct.pack('>h6s', 6, b"member"),  # Member ID
            struct.pack('>i', 2),             # Two member datas

            struct.pack('>h3s', 3, b"id1"),     # ID
            struct.pack('>i5s', 5, b"data1"),   # Metadata
            struct.pack('>h3s', 3, b"id2"),     # ID
            struct.pack('>i5s', 5, b"data2"),   # Metadata
        ])
        decoded = KafkaCodec.decode_join_group_response(encoded)
        self.assertEqual(decoded, expected)

    def test_encode_heartbeat_request(self):
        expected = b"".join([
            struct.pack('>h', 12),          # API key HeartbeatRequest
            struct.pack('>h', 0),           # API version
            struct.pack('>i', 1),           # Correlation ID
            struct.pack('>h3s', 3, b"cID"),      # The client ID
            struct.pack('>h6s', 6, b"group1"),   # Group
            struct.pack('>i', 1),                # generation ID
            struct.pack('>h6s', 6, b"member"),   # Member ID
        ])

        encoded = KafkaCodec.encode_heartbeat_request(
            b"cID", 1,
            _HeartbeatRequest("group1", 1, "member"),
        )

        self.assertEqual(encoded, expected)

    def test_decode_heartbeat_response(self):
        expected = _HeartbeatResponse(0)

        encoded = b"".join([
            struct.pack('>i', 9),  # Correlation ID
            struct.pack('>h', 0),  # Error Code
        ])
        decoded = KafkaCodec.decode_heartbeat_response(encoded)
        self.assertEqual(decoded, expected)

    def test_encode_sync_group_request(self):
        expected = b"".join([
            struct.pack('>h', 14),               # API key SyncGroupRequest
            struct.pack('>h', 0),                # API version
            struct.pack('>i', 1),                # Correlation ID
            struct.pack('>h3s', 3, b"cID"),      # The client ID
            struct.pack('>h6s', 6, b"group1"),   # Consumer group 'group1'
            struct.pack('>i', 1),                # Generation ID
            struct.pack('>h6s', 6, b"member"),   # Member ID
            struct.pack('>i', 2),                # Two group assignments

            struct.pack('>h4s', 4, b"name"),     # Name
            struct.pack('>i4s', 4, b"meta"),     # Metadata
            struct.pack('>h5s', 5, b"name2"),    # Name
            struct.pack('>i5s', 5, b"meta2"),    # Metadata
        ])

        encoded = KafkaCodec.encode_sync_group_request(
            b"cID", 1,
            _SyncGroupRequest("group1", 1, "member", [
                _SyncGroupRequestMember("name", b"meta"),
                _SyncGroupRequestMember("name2", b"meta2"),
            ]),
        )
        self.assertEqual(encoded, expected)

    def test_decode_sync_group_response(self):
        expected = _SyncGroupResponse(0, b"data")

        encoded = b"".join([
            struct.pack('>i', 9),  # Correlation ID
            struct.pack('>h', 0),  # Error Code
            struct.pack('>i4s', 4, b"data"),   # Metadata
        ])
        decoded = KafkaCodec.decode_sync_group_response(encoded)
        self.assertEqual(decoded, expected)

    def test_encode_sync_group_member_assignment(self):
        expected = b"".join([
            struct.pack('>h', 1),               # version 1
            struct.pack('>i', 1),               # one assignment
            struct.pack('>h5s', 5, b"topic"),   # topic
            struct.pack('>ii', 1, 5),           # one partition: 5
            struct.pack('>i4s', 4, b"data"),    # Metadata
        ])

        encoded = KafkaCodec.encode_sync_group_member_assignment(
            1,
            {"topic": [5]},
            b"data",
        )
        self.assertEqual(encoded, expected)

    def test_decode_sync_group_member_assignment(self):
        expected = _SyncGroupMemberAssignment(0, {"topic": (5,)}, b"data")

        encoded = b"".join([
            struct.pack('>h', 0),               # version 0
            struct.pack('>i', 1),               # one assignment
            struct.pack('>h5s', 5, b"topic"),   # topic
            struct.pack('>ii', 1, 5),           # one partition: 5
            struct.pack('>i4s', 4, b"data"),    # Metadata
        ])
        decoded = KafkaCodec.decode_sync_group_member_assignment(encoded)
        self.assertEqual(decoded, expected)

    def test_decode_sync_group_member_assignment_bad_version(self):
        """
        Afkak only produces a version 0 SyncGroupMemberAssignment, so any other
        version produces a ProtocolError.
        """
        encoded = b"".join([
            struct.pack('>h', 1),               # version 1
            struct.pack('>i', 1),               # one assignment
            struct.pack('>h5s', 5, b"topic"),   # topic
            struct.pack('>ii', 1, 5),           # one partition: 5
            struct.pack('>i4s', 4, b"data"),    # Metadata
        ])
        self.assertRaises(ProtocolError, KafkaCodec.decode_sync_group_member_assignment, encoded)

    def test_encode_leave_group_request(self):
        expected = b"".join([
            struct.pack('>h', 13),               # API key SyncGroupRequest
            struct.pack('>h', 0),                # API version
            struct.pack('>i', 1),                # Correlation ID
            struct.pack('>h3s', 3, b"cID"),      # The client ID
            struct.pack('>h6s', 6, b"group1"),   # Consumer group 'group1'
            struct.pack('>h6s', 6, b"member"),   # Member ID
        ])

        encoded = KafkaCodec.encode_leave_group_request(
            b"cID", 1,
            _LeaveGroupRequest("group1", "member"),
        )
        self.assertEqual(encoded, expected)

    def test_decode_leave_group_response(self):
        expected = _LeaveGroupResponse(0)

        encoded = struct.pack('>ih', 9, 0)  # Correlation ID, error code
        decoded = KafkaCodec.decode_leave_group_response(encoded)
        self.assertEqual(decoded, expected)
