/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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.
 */
package org.apache.kafka.clients.consumer.internals.events;

import org.apache.kafka.clients.consumer.internals.ConsumerUtils;
import org.apache.kafka.common.errors.TimeoutException;
import org.apache.kafka.common.utils.LogContext;
import org.apache.kafka.common.utils.MockTime;
import org.apache.kafka.common.utils.Time;
import org.junit.jupiter.api.Test;

import java.util.ArrayList;
import java.util.Collection;

import static org.apache.kafka.clients.consumer.internals.events.CompletableEvent.calculateDeadlineMs;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;

public class CompletableEventReaperTest {

    private final LogContext logContext = new LogContext();
    private final Time time = new MockTime();
    private final CompletableEventReaper reaper = new CompletableEventReaper(logContext);

    @Test
    public void testExpired() {
        // Add a new event to the reaper.
        long timeoutMs = 100;
        UnsubscribeEvent event = new UnsubscribeEvent(calculateDeadlineMs(time.milliseconds(), timeoutMs));
        reaper.add(event);

        // Without any time passing, we check the reaper and verify that the event is not done amd is still
        // being tracked.
        reaper.reap(time.milliseconds());
        assertFalse(event.future().isDone());
        assertEquals(1, reaper.size());

        // Sleep for at least 1 ms. *more* than the timeout so that the event is considered expired.
        time.sleep(timeoutMs + 1);

        // However, until we actually invoke the reaper, the event isn't complete and is still being tracked.
        assertFalse(event.future().isDone());
        assertEquals(1, reaper.size());

        // Call the reaper and validate that the event is now "done" (expired), the correct exception type is
        // thrown, and the event is no longer tracked.
        reaper.reap(time.milliseconds());
        assertTrue(event.future().isDone());
        assertThrows(TimeoutException.class, () -> ConsumerUtils.getResult(event.future()));
        assertEquals(0, reaper.size());
    }

    @Test
    public void testCompleted() {
        // Add a new event to the reaper.
        long timeoutMs = 100;
        UnsubscribeEvent event = new UnsubscribeEvent(calculateDeadlineMs(time.milliseconds(), timeoutMs));
        reaper.add(event);

        // Without any time passing, we check the reaper and verify that the event is not done amd is still
        // being tracked.
        reaper.reap(time.milliseconds());
        assertFalse(event.future().isDone());
        assertEquals(1, reaper.size());

        // We'll cause the event to be completed normally. Note that because we haven't called the reaper, the
        // event is still being tracked.
        event.future().complete(null);
        assertTrue(event.future().isDone());
        assertEquals(1, reaper.size());

        // To ensure we don't accidentally expire an event that completed normally, sleep past the timeout.
        time.sleep(timeoutMs + 1);

        // Call the reaper and validate that the event is not considered expired, but is still no longer tracked.
        reaper.reap(time.milliseconds());
        assertTrue(event.future().isDone());
        assertNull(ConsumerUtils.getResult(event.future()));
        assertEquals(0, reaper.size());
    }

    @Test
    public void testCompletedAndExpired() {
        // Add two events to the reaper. One event will be completed, the other we will let expire.
        long timeoutMs = 100;
        UnsubscribeEvent event1 = new UnsubscribeEvent(calculateDeadlineMs(time.milliseconds(), timeoutMs));
        UnsubscribeEvent event2 = new UnsubscribeEvent(calculateDeadlineMs(time.milliseconds(), timeoutMs));
        reaper.add(event1);
        reaper.add(event2);

        // Without any time passing, we check the reaper and verify that the event is not done amd is still
        // being tracked.
        reaper.reap(time.milliseconds());
        assertFalse(event1.future().isDone());
        assertFalse(event2.future().isDone());
        assertEquals(2, reaper.size());

        // We'll cause the first event to be completed normally, but then sleep past the timer deadline.
        event1.future().complete(null);
        assertTrue(event1.future().isDone());

        time.sleep(timeoutMs + 1);

        // Though the first event is completed, it's still being tracked, along with the second expired event.
        assertEquals(2, reaper.size());

        // Validate that the first (completed) event is not expired, but the second one is expired. In either case,
        // both should be completed and neither should be tracked anymore.
        reaper.reap(time.milliseconds());
        assertTrue(event1.future().isDone());
        assertTrue(event2.future().isDone());
        assertNull(ConsumerUtils.getResult(event1.future()));
        assertThrows(TimeoutException.class, () -> ConsumerUtils.getResult(event2.future()));
        assertEquals(0, reaper.size());
    }

    @Test
    public void testIncompleteQueue() {
        long timeoutMs = 100;
        UnsubscribeEvent event1 = new UnsubscribeEvent(calculateDeadlineMs(time.milliseconds(), timeoutMs));
        UnsubscribeEvent event2 = new UnsubscribeEvent(calculateDeadlineMs(time.milliseconds(), timeoutMs));

        Collection<CompletableApplicationEvent<?>> queue = new ArrayList<>();
        queue.add(event1);
        queue.add(event2);

        // Complete one of our events, just to make sure it isn't inadvertently canceled.
        event1.future().complete(null);

        // In this test, our events aren't tracked in the reaper, just in the queue.
        assertEquals(0, reaper.size());
        assertEquals(2, queue.size());

        // Go ahead and reap the incomplete from the queue.
        reaper.reap(queue);

        // The first event was completed, so we didn't expire it in the reaper.
        assertTrue(event1.future().isDone());
        assertFalse(event1.future().isCancelled());
        assertNull(ConsumerUtils.getResult(event1.future()));

        // The second event was incomplete, so it was expired.
        assertTrue(event2.future().isCompletedExceptionally());
        assertThrows(TimeoutException.class, () -> ConsumerUtils.getResult(event2.future()));

        // Because the events aren't tracked in the reaper *and* the queue is cleared as part of the
        // cancellation process, our data structures should both be the same as above.
        assertEquals(0, reaper.size());
        assertEquals(0, queue.size());
    }

    @Test
    public void testIncompleteTracked() {
        // This queue is just here to test the case where the queue is empty.
        Collection<CompletableApplicationEvent<?>> queue = new ArrayList<>();

        // Add two events for the reaper to track.
        long timeoutMs = 100;
        UnsubscribeEvent event1 = new UnsubscribeEvent(calculateDeadlineMs(time.milliseconds(), timeoutMs));
        UnsubscribeEvent event2 = new UnsubscribeEvent(calculateDeadlineMs(time.milliseconds(), timeoutMs));
        reaper.add(event1);
        reaper.add(event2);

        // Complete one of our events, just to make sure it isn't inadvertently canceled.
        event1.future().complete(null);

        // In this test, our events are tracked exclusively in the reaper, not the queue.
        assertEquals(2, reaper.size());

        // Go ahead and reap the incomplete events. Both sets should be zero after that.
        reaper.reap(queue);
        assertEquals(0, reaper.size());
        assertEquals(0, queue.size());

        // The first event was completed, so we didn't cancel it in the reaper.
        assertTrue(event1.future().isDone());
        assertNull(ConsumerUtils.getResult(event1.future()));

        // The second event was incomplete, so it was canceled.
        assertTrue(event2.future().isCompletedExceptionally());
        assertThrows(TimeoutException.class, () -> ConsumerUtils.getResult(event2.future()));
    }
}
