From 8eaf8b6dd985c5d15f2433e15641781dc0ca7317 Mon Sep 17 00:00:00 2001
From: Artem Bilan <abilan@vmware.com>
Date: Thu, 25 Mar 2021 10:34:35 -0400
Subject: [PATCH] Make CallerBlocksPolicyTests more robust

https://build.spring.io/browse/INT-MASTER-2372/
---
 .../util/CallerBlocksPolicyTests.java         | 82 +++++++++++--------
 1 file changed, 47 insertions(+), 35 deletions(-)

diff --git a/spring-integration-core/src/test/java/org/springframework/integration/util/CallerBlocksPolicyTests.java b/spring-integration-core/src/test/java/org/springframework/integration/util/CallerBlocksPolicyTests.java
index 864e8312fc..1a74432b3a 100644
--- a/spring-integration-core/src/test/java/org/springframework/integration/util/CallerBlocksPolicyTests.java
+++ b/spring-integration-core/src/test/java/org/springframework/integration/util/CallerBlocksPolicyTests.java
@@ -23,7 +23,7 @@ import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
 
-import org.junit.Test;
+import org.junit.jupiter.api.Test;
 
 import org.springframework.core.task.TaskRejectedException;
 import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
@@ -39,14 +39,14 @@ public class CallerBlocksPolicyTests {
 
 	@Test
 	public void test0() throws Exception {
-		final ThreadPoolTaskExecutor te = new ThreadPoolTaskExecutor();
+		ThreadPoolTaskExecutor te = new ThreadPoolTaskExecutor();
 		te.setCorePoolSize(1);
 		te.setMaxPoolSize(1);
 		te.setQueueCapacity(0);
 		te.setRejectedExecutionHandler(new CallerBlocksPolicy(10));
 		te.initialize();
-		final AtomicReference<Throwable> e = new AtomicReference<>();
-		final CountDownLatch latch = new CountDownLatch(1);
+		AtomicReference<Throwable> e = new AtomicReference<>();
+		CountDownLatch latch = new CountDownLatch(1);
 		Runnable task = new Runnable() {
 
 			@Override
@@ -60,46 +60,58 @@ public class CallerBlocksPolicyTests {
 				latch.countDown();
 			}
 		};
-		te.execute(task);
-		assertThat(latch.await(10, TimeUnit.SECONDS)).isTrue();
-		assertThat(e.get()).isInstanceOf(RejectedExecutionException.class);
-		assertThat(e.get().getMessage()).isEqualTo("Max wait time expired to queue task");
-		te.destroy();
+		try {
+			te.execute(task);
+			assertThat(latch.await(10, TimeUnit.SECONDS)).isTrue();
+			assertThat(e.get())
+					.isInstanceOf(RejectedExecutionException.class)
+					.hasMessage("Max wait time expired to queue task");
+		}
+		finally {
+			te.destroy();
+		}
 	}
 
 	@Test
 	public void test1() throws Exception {
-		final ThreadPoolTaskExecutor te = new ThreadPoolTaskExecutor();
+		ThreadPoolTaskExecutor te = new ThreadPoolTaskExecutor();
 		te.setCorePoolSize(2);
 		te.setMaxPoolSize(2);
 		te.setQueueCapacity(1);
 		te.setRejectedExecutionHandler(new CallerBlocksPolicy(10000));
 		te.initialize();
-		final AtomicReference<Throwable> e = new AtomicReference<Throwable>();
-		final CountDownLatch latch = new CountDownLatch(3);
-		te.execute(() -> {
-			try {
-				Runnable foo = () -> {
-					try {
-						Thread.sleep(10);
-					}
-					catch (InterruptedException e1) {
-						Thread.currentThread().interrupt();
-						throw new RuntimeException();
-					}
-					latch.countDown();
-				};
-				te.execute(foo);
-				te.execute(foo); // this one will be queued
-				te.execute(foo); // this one will be blocked and successful later
-			}
-			catch (TaskRejectedException tre) {
-				e.set(tre.getCause());
-			}
-		});
-		assertThat(latch.await(10, TimeUnit.SECONDS)).isTrue();
-		assertThat(e.get()).isNull();
-		te.destroy();
+		AtomicReference<Throwable> e = new AtomicReference<>();
+		CountDownLatch latch = new CountDownLatch(3);
+		try {
+			te.execute(() -> {
+				try {
+					Runnable foo =
+							() -> {
+								try {
+									Thread.sleep(10);
+								}
+								catch (InterruptedException e1) {
+									Thread.currentThread().interrupt();
+									throw new RuntimeException(e1);
+								}
+								finally {
+									latch.countDown();
+								}
+							};
+					te.execute(foo);
+					te.execute(foo); // this one will be queued
+					te.execute(foo); // this one will be blocked and successful later
+				}
+				catch (TaskRejectedException tre) {
+					e.set(tre.getCause());
+				}
+			});
+			assertThat(latch.await(10, TimeUnit.SECONDS)).isTrue();
+			assertThat(e.get()).isNull();
+		}
+		finally {
+			te.destroy();
+		}
 	}
 
 }