"""Test Kafka helpers."""
from kafka.consumer.fetcher import ConsumerRecord
from kafka.producer.future import RecordMetadata
from kafka.structs import OffsetAndMetadata
from kafka.structs import TopicPartition
import pytest

from dbdeploy.util.exceptions import KafkaSendError
from dbdeploy.util.kafka.structs import PartitionOffset
from dbdeploy.util.kafka import helpers


class TestKafkaProducerCallbacks:
    def test_on_send_success_callback(self):
        """Test execute callback on successfully sent messages."""
        record_metadata = RecordMetadata(
            topic='test_topic',
            partition=1,
            topic_partition=TopicPartition(topic='test_topic', partition=1),
            offset=1,
            timestamp='2022-01-01',
            log_start_offset=1,
            checksum=1,
            serialized_header_size=1,
            serialized_key_size=1,
            serialized_value_size=1,
        )
        latest_offsets = {}
        helpers.on_send_success_callback(
            record_metadata=record_metadata,
            latest_offsets=latest_offsets)
        assert {1: 1} == latest_offsets

    def test_on_send_error_callback(self):
        """Test execute callback on message sending error."""
        with pytest.raises(KafkaSendError):
            helpers.on_send_error_callback('test_err')


class TestOffests:
    def test_get_offsets_per_partiton(self):
        """Test return a list of consumer's topic-offest tuples."""
        topic_partition = TopicPartition('test_topic', 1)
        offset_and_metadata = OffsetAndMetadata(offset='1', metadata=None)
        test_result = helpers.get_offsets_per_partiton(
            topic='test_topic',
            offsets={topic_partition: offset_and_metadata})
        assert test_result == [PartitionOffset(partition='1', offset='1')]

    def test_compare_offsets(self):
        """Test latest offset dict gets updated."""
        latest_offsets = {1: 1}
        current_offsets = [PartitionOffset(1, 1)]
        helpers.compare_offsets(
            latest_offsets=latest_offsets,
            current_offsets=current_offsets)
        assert latest_offsets == {}

    def test_compare_offsets_latest_offset_is_not_reached(self):
        """Test latest offset dict isn't updated."""
        latest_offsets = {1: 5}
        current_offsets = [PartitionOffset(1, 1)]
        helpers.compare_offsets(
            latest_offsets=latest_offsets,
            current_offsets=current_offsets)
        assert latest_offsets == {1: 5}

    def test_compare_offsets_only_one_partition_popped(self):
        """Test latest offset dict isn't updated."""
        latest_offsets = {1: 5, 2: 2}
        current_offsets = [PartitionOffset(2, 12), PartitionOffset(1, 1)]
        helpers.compare_offsets(
            latest_offsets=latest_offsets,
            current_offsets=current_offsets)
        assert latest_offsets == {1: 5}


class TestKafkaDLQHelpers:
    def test_check_dlq_topic_records(self):
        """Test DLQ check returns exception."""
        topic_partition = TopicPartition('test_topic', 1)
        record = ConsumerRecord(
            'test_topic', 1, 1, 'timestamp',
            'timestamp_type', b'test_key', b'test_value', 'headers',
            'checksum', 'serialized_key_size', 'serialized_value_size',
            'serialized_header_size')
        polled_records = {topic_partition: [record]}
        partitions = [topic_partition]
        test_result = helpers.check_dlq_topic_records(
                    polled_records=polled_records,
                    partitions=partitions)
        assert test_result == {'stop': True, 'record': record}