add tests

jrhee17 · Sep 26, 2023 · afcd9d2 · afcd9d2
1 parent af7537d
commit afcd9d2
Showing 1 changed file with 56 additions and 0 deletions.
diff --git a/core/src/test/java/com/linecorp/armeria/internal/common/CancellationSchedulerTest.java b/core/src/test/java/com/linecorp/armeria/internal/common/CancellationSchedulerTest.java
@@ -28,11 +28,14 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
+import java.util.stream.Stream;
 
 import org.assertj.core.data.Offset;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.CsvSource;
+import org.junit.jupiter.params.provider.ValueSource;
 
 import com.linecorp.armeria.client.ResponseTimeoutException;
 import com.linecorp.armeria.common.CommonPools;
@@ -57,6 +60,21 @@ public boolean canSchedule() {
         public void run(Throwable cause) {}
     };
 
+    static final class StatefulCancellationTask implements CancellationTask {
+
+        AtomicReference<Throwable> thrownRef = new AtomicReference<>();
+
+        @Override
+        public boolean canSchedule() {
+            return true;
+        }
+
+        @Override
+        public void run(Throwable cause) {
+            thrownRef.set(cause);
+        }
+    }
+
     private static void executeInEventLoop(long initTimeoutNanos, Consumer<CancellationScheduler> task) {
         final AtomicBoolean completed = new AtomicBoolean();
         eventExecutor.execute(() -> {
@@ -441,6 +459,44 @@ void multiple_ClearTimeoutInWhenCancelling() {
         await().untilTrue(completed);
     }
 
+    @Test
+    void immediateFinishTriggersCompletion() {
+        final CancellationScheduler scheduler = new CancellationScheduler(0);
+        final Throwable throwable = new Throwable();
+
+        assertThat(scheduler.whenCancelling()).isNotCompleted();
+        assertThat(scheduler.state()).isEqualTo(State.INIT);
+
+        scheduler.finishNow(throwable);
+
+        assertThat(scheduler.whenCancelling()).isNotDone();
+        assertThat(scheduler.whenCancelled()).isCompleted();
+        assertThat(scheduler.state()).isEqualTo(State.FINISHED);
+        assertThat(scheduler.cause()).isSameAs(throwable);
+    }
+
+    @Test
+    void finishAfterInitNotComplete() {
+        final CancellationScheduler scheduler = new CancellationScheduler(0);
+        final Throwable throwable = new Throwable();
+
+        assertThat(scheduler.whenCancelling()).isNotCompleted();
+        assertThat(scheduler.state()).isEqualTo(State.INIT);
+
+        final StatefulCancellationTask task = new StatefulCancellationTask();
+        scheduler.init(eventExecutor, task, 0, false);
+        scheduler.finishNow(throwable);
+
+        await().untilAsserted(() -> assertThat(scheduler.whenCancelled()).isDone());
+
+        assertThat(scheduler.whenCancelled()).isCompleted();
+        assertThat(scheduler.state()).isEqualTo(State.FINISHED);
+
+        // verify the task was actually executed
+        assertThat(scheduler.whenCancelling()).isCompleted();
+        assertThat(task.thrownRef.get()).isSameAs(throwable);
+    }
+
     static void assertTimeoutWithTolerance(long actualNanos, long expectedNanos) {
         assertTimeoutWithTolerance(actualNanos, expectedNanos, MILLISECONDS.toNanos(200));
     }