"""Testing the KafkaExecutor."""
from collections import namedtuple

from kafka.structs import OffsetAndMetadata
from kafka.structs import TopicPartition
import pytest

from feed_ingestion.util.kafka import helpers
from feed_ingestion.util.kafka.structs import PartitonOffset


Record = namedtuple('Record', ['offset', 'key'])
"""Mock topic record."""


@pytest.mark.parametrize(
    'topic, offsets, result',
    [
        ('test-topic', {0: 11, 1: 25, 2: 22}, True),
        ('test-topic', {0: 51, 1: 15, 2: 87}, True),
        ('false-topic', {0: 41, 1: 48, 2: 38}, False),
        ('test-topic', {0: 31, 1: 15, 2: 28}, True),
        ('test-topic', {0: 81, 1: 95, 2: 82}, True),
    ]
)
def test_get_offsets_per_partiton(topic, offsets, result):
    """Test helpers returns list of partitions."""
    consumer_offsets = {}
    # generate an example of what normally
    # returned by boto3.client.list_offsest()
    for k, v in offsets.items():
        tp = TopicPartition('test-topic', k)
        consumer_offsets[tp] = OffsetAndMetadata(v, 'meta')
    offset_dict = helpers.get_offsets_per_partiton(topic, consumer_offsets)
    assert offset_dict == [
        PartitonOffset(partition='0', offset=offsets[0]),
        PartitonOffset(partition='1', offset=offsets[1]),
        PartitonOffset(partition='2', offset=offsets[2]),
    ] if result else offset_dict == []


@pytest.mark.parametrize(
    'latest_offsets, current_offsets, finished_len, unfinished_len',
    [
        ({0: 11, 1: 25, 2: 22}, [
            PartitonOffset('0', 5),
            PartitonOffset('1', 6),
            PartitonOffset('2', 7),
        ], 0, 3),
        ({0: 11, 1: 25, 2: 22}, [
            PartitonOffset('0', 55),
            PartitonOffset('1', 6),
            PartitonOffset('2', 7),
        ], 1, 2),
        ({0: 11, 1: 25, 2: 22}, [
            PartitonOffset('0', 55),
            PartitonOffset('1', 6),
            PartitonOffset('2', 77),
        ], 2, 1),
        ({0: 11, 1: 25, 2: 22}, [
            PartitonOffset('0', 55),
            PartitonOffset('1', 62),
            PartitonOffset('2', 71),
        ], 3, 0),
        ({0: 0, 1: 0, 2: 1}, [
            PartitonOffset('0', 0),
            PartitonOffset('1', 0),
            PartitonOffset('2', 1),
        ], 3, 0),
    ]
)
def test_compare_current_vs_latest_offsets(
        latest_offsets, current_offsets, finished_len, unfinished_len):
    """Compare consumer offsets and update finished/unfinished sets."""
    finished_partitions = set()
    unfinished_partitions = {1, 2, 3}
    helpers.compare_current_vs_latest_offsets(
        latest_offsets=latest_offsets,
        current_offsets=current_offsets,
        finished_partitions=finished_partitions,
        unfinished_partitions=unfinished_partitions)
    expected = (finished_len, unfinished_len)
    assert (len(finished_partitions), len(unfinished_partitions)) == expected


@pytest.mark.parametrize(
    'sf_query_id, offset, records, stop_result',
    [
        ('test-query-id', 11,
         [Record(offset, b'{"m":"n","key":"false"}')
          for offset in range(50)], False),
        ('test-query-id', 500,
         [Record(121, b'{"a":"b","key":"test-query-id"}')], True),
        ('test-query-id', 1234,
         [Record(1, b'{"c":"d","key":"false"}'),
          Record(1001, b'{"k":"test-query-id"}')], True),
        ('test-query-id', 1234,
         [Record(offset, b'{"c":"d","key":"false"}')
          for offset in range(1000)], True),
        ('test-query-id', 951,
         [Record(offset, b'{"e":"f","key":"false"}')
          for offset in range(1000)], False),
        ('test-query-id', 101,
         [Record(10001, b'{"j":"k","key":"test-query-id"}')], False),
    ]
)
def test_check_dlq_topic_records(
        sf_query_id, offset, records, stop_result):
    """Look for exeptions in DLQ."""
    tp = TopicPartition('test-dlq-topic', 0)
    if stop_result:
        record = Record(1001, b'{"k":"test-query-id"}')
        records.append(record)
    end_offsets = {tp: offset}
    polled_records = {tp: records}
    result = helpers.check_dlq_topic_records(
        sf_query_id,
        polled_records,
        [tp],
        end_offsets)
    if stop_result:
        assert result.get('stop')
    else:
        assert not result.get('stop')