Fix flaky tests

line · May 23, 2024 · 51caf2a · 51caf2a
1 parent 806556e
commit 51caf2a
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/core/src/test/java/com/linecorp/armeria/common/stream/ByteStreamMessageOutputStreamTest.java b/core/src/test/java/com/linecorp/armeria/common/stream/ByteStreamMessageOutputStreamTest.java
@@ -29,6 +29,7 @@
 import java.util.concurrent.atomic.AtomicInteger;
 
 import org.hamcrest.Matchers;
+import org.junit.jupiter.api.RepeatedTest;
 import org.junit.jupiter.api.Test;
 
 import com.linecorp.armeria.common.HttpData;
@@ -152,6 +153,7 @@ void writeAfterClosed() {
 
     @Test
     void writeAfterStreamClosed() throws InterruptedException {
+        final CountDownLatch abortWait = new CountDownLatch(1);
         final CountDownLatch wait = new CountDownLatch(1);
         final CountDownLatch end = new CountDownLatch(1);
         final ByteStreamMessage byteStreamMessage = StreamMessage.fromOutputStream(os -> {
@@ -160,6 +162,7 @@ void writeAfterStreamClosed() throws InterruptedException {
                     if (i < 2) {
                         os.write(i);
                     } else {
+                        abortWait.countDown();
                         wait.await();
                         assertThatThrownBy(() -> os.write(0))
                                 .isInstanceOf(IOException.class)
@@ -175,6 +178,13 @@ void writeAfterStreamClosed() throws InterruptedException {
 
         StepVerifier.create(byteStreamMessage, 2)
                     .expectNext(httpData(0), httpData(1))
+                    .then(() -> {
+                        try {
+                            abortWait.await();
+                        } catch (InterruptedException e) {
+                            throw new RuntimeException(e);
+                        }
+                    })
                     .then(byteStreamMessage::abort)
                     .then(() -> {
                         // Wait for the abortion to be completed.

diff --git a/core/src/test/java12/com/linecorp/armeria/server/InvalidPathWithDataTest.java b/core/src/test/java12/com/linecorp/armeria/server/InvalidPathWithDataTest.java
@@ -27,9 +27,11 @@
 
 import com.linecorp.armeria.common.HttpResponse;
 import com.linecorp.armeria.common.RequestTarget;
+import com.linecorp.armeria.internal.testing.FlakyTest;
 import com.linecorp.armeria.server.logging.LoggingService;
 import com.linecorp.armeria.testing.junit5.server.ServerExtension;
 
+@FlakyTest
 class InvalidPathWithDataTest {
 
     @RegisterExtension