package com.lingniu.ingest.sink.kafka;

import org.junit.jupiter.api.Test;

import java.time.Duration;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.mock;

class KafkaEnvelopeConsumerRunnerTest {

    @Test
    void stopWakesPollLoopAndClosesWorkerAfterPollThreadExits() throws Exception {
        KafkaEnvelopeConsumerWorker worker = mock(KafkaEnvelopeConsumerWorker.class);
        CountDownLatch pollEntered = new CountDownLatch(1);
        CountDownLatch wakeupCalled = new CountDownLatch(1);
        CountDownLatch pollReturned = new CountDownLatch(1);
        CountDownLatch closeCalled = new CountDownLatch(1);

        doAnswer(invocation -> {
            pollEntered.countDown();
            assertThat(wakeupCalled.await(1, TimeUnit.SECONDS)).isTrue();
            pollReturned.countDown();
            return 0;
        }).when(worker).pollOnce(any());
        doAnswer(invocation -> {
            wakeupCalled.countDown();
            return null;
        }).when(worker).wakeup();
        doAnswer(invocation -> {
            assertThat(pollReturned.getCount()).isZero();
            closeCalled.countDown();
            return null;
        }).when(worker).close();
        KafkaEnvelopeConsumerRunner runner = new KafkaEnvelopeConsumerRunner(
                List.of(worker),
                Duration.ofSeconds(30),
                Duration.ofMillis(1),
                false);

        runner.start();
        assertThat(pollEntered.await(1, TimeUnit.SECONDS)).isTrue();
        runner.stop();

        assertThat(wakeupCalled.getCount()).isZero();
        assertThat(closeCalled.getCount()).isZero();
    }
}
