package com.criteo.hadoop.garmadon.hdfs.kafka;

import com.criteo.hadoop.garmadon.hdfs.offset.OffsetComputer;
import com.criteo.hadoop.garmadon.hdfs.writer.PartitionedWriter;
import org.apache.kafka.clients.consumer.Consumer;
import org.apache.kafka.common.TopicPartition;
import org.junit.Test;

import java.io.IOException;
import java.util.*;

import static org.mockito.Mockito.*;

@SuppressWarnings("unchecked")
public class OffsetResetterTest {
    private static final String TOPIC = "topic";

    @Test
    public void partitionAssignedValidOffset() throws IOException {
        final Consumer<Long, String> consumer = mock(Consumer.class);
        final PartitionedWriter firstWriter = mock(PartitionedWriter.class);
        final PartitionedWriter secondWriter = mock(PartitionedWriter.class);
        final PartitionedWriter newEventWriter = mock(PartitionedWriter.class);
        final OffsetResetter offsetResetter = new OffsetResetter<>(consumer, mock(java.util.function.Consumer.class),
            Arrays.asList(firstWriter, secondWriter, newEventWriter));
        final TopicPartition firstPartition = new TopicPartition(TOPIC, 1);
        final TopicPartition secondPartition = new TopicPartition(TOPIC, 2);
        final List<TopicPartition> partitions = Arrays.asList(firstPartition, secondPartition);

        final Map<Integer, Long> firstOffsets = new HashMap<>();
        firstOffsets.put(firstPartition.partition(), 10L);
        firstOffsets.put(secondPartition.partition(), OffsetComputer.NO_OFFSET);
        when(firstWriter.getStartingOffsets(any())).thenReturn(firstOffsets);

        final Map<Integer, Long> secondOffsets = new HashMap<>();
        secondOffsets.put(firstPartition.partition(), 15L);
        secondOffsets.put(secondPartition.partition(), OffsetComputer.NO_OFFSET);
        when(secondWriter.getStartingOffsets(any())).thenReturn(secondOffsets);

        final Map<Integer, Long> thirdOffsets = new HashMap<>();
        thirdOffsets.put(firstPartition.partition(), OffsetComputer.NO_OFFSET);
        thirdOffsets.put(secondPartition.partition(), OffsetComputer.NO_OFFSET);
        when(newEventWriter.getStartingOffsets(any())).thenReturn(thirdOffsets);

        offsetResetter.onPartitionsAssigned(partitions);
        verify(consumer, times(1)).seek(eq(firstPartition), eq(10L));
        verify(consumer, times(1)).seekToBeginning(Collections.singleton(secondPartition));
        verifyNoMoreInteractions(consumer);
    }

    @Test
    public void partitionsAssignedCannotFetchOffset() throws IOException {
        final Consumer<Long, String> consumer = mock(Consumer.class);
        final PartitionedWriter successfulWriter = mock(PartitionedWriter.class);
        final PartitionedWriter exceptionalWriter = mock(PartitionedWriter.class);
        final OffsetResetter offsetResetter = new OffsetResetter<>(consumer, mock(java.util.function.Consumer.class),
                Arrays.asList(successfulWriter, exceptionalWriter));
        final TopicPartition partition = new TopicPartition(TOPIC, 1);
        final List<TopicPartition> partitions = Collections.singletonList(partition);

        when(successfulWriter.getStartingOffsets(any())).thenReturn(new HashMap<>());
        when(exceptionalWriter.getStartingOffsets(any())).thenThrow(new IOException("Ayo"));

        offsetResetter.onPartitionsAssigned(partitions);
        verify(consumer, times(1)).seekToBeginning(Collections.singleton(partition));
        verifyNoMoreInteractions(consumer);
    }

    @Test
    public void partitionsAssignedNoWriter() {
        final Consumer<Long, String> consumer = mock(Consumer.class);
        final OffsetResetter offsetResetter = new OffsetResetter<>(consumer, mock(java.util.function.Consumer.class),
                Collections.emptyList());
        final TopicPartition partition = new TopicPartition(TOPIC, 1);
        final List<TopicPartition> partitions = Collections.singletonList(partition);

        offsetResetter.onPartitionsAssigned(partitions);
        verify(consumer, times(1)).seekToBeginning(Collections.singleton(partition));
        verifyNoMoreInteractions(consumer);
    }

    @Test
    public void partitionsAssignedNoTopic() {
        final Consumer<Long, String> consumer = mock(Consumer.class);
        final OffsetResetter offsetResetter = new OffsetResetter<>(consumer, mock(java.util.function.Consumer.class),
                Collections.singleton(mock(PartitionedWriter.class)));

        offsetResetter.onPartitionsAssigned(Collections.emptyList());
        verifyZeroInteractions(consumer);
    }

    @Test
    public void partitionsRevoked() {
        final Consumer<Long, String> consumer = mock(Consumer.class);
        final PartitionedWriter writer = mock(PartitionedWriter.class);
        final java.util.function.Consumer partitionsRevokedConsumer = mock(java.util.function.Consumer.class);
        final OffsetResetter offsetResetter = new OffsetResetter<>(consumer, partitionsRevokedConsumer,
                Collections.singleton(writer));

        offsetResetter.onPartitionsRevoked(Arrays.asList(new TopicPartition(TOPIC, 1),
                new TopicPartition(TOPIC, 2)));
        verify(writer, times(1)).close();
        verify(writer, times(1)).dropPartition(1);
        verify(writer, times(1)).dropPartition(2);
        verify(partitionsRevokedConsumer, times(1)).accept(1);
        verify(partitionsRevokedConsumer, times(1)).accept(2);
        verifyNoMoreInteractions(writer);
    }

    @Test
    public void partitionsRevokedNoTopic() {
        final Consumer<Long, String> consumer = mock(Consumer.class);
        final PartitionedWriter writer = mock(PartitionedWriter.class);
        final java.util.function.Consumer partitionsRevokedConsumer = mock(java.util.function.Consumer.class);
        final OffsetResetter offsetResetter = new OffsetResetter<>(consumer, partitionsRevokedConsumer,
                Collections.singleton(writer));

        offsetResetter.onPartitionsRevoked(Collections.emptyList());
        verify(writer, times(1)).close();
        verifyZeroInteractions(partitionsRevokedConsumer);
    }

    @Test
    public void partitionsRevokedNoWriter() {
        final Consumer<Long, String> consumer = mock(Consumer.class);
        final OffsetResetter offsetResetter = new OffsetResetter<>(consumer, mock(java.util.function.Consumer.class),
                Collections.emptyList());

        offsetResetter.onPartitionsRevoked(Collections.singleton(new TopicPartition(TOPIC, 1)));
    }
}