ExpectedSystemExit improvement in multi-threads environment

src/main/java/org/junit/contrib/java/lang/system/ExpectedSystemExit.java

@@ -28,7 +28,9 @@
  * <p>
  * Some care must be taken if your system under test creates a new thread and
  * this thread calls {@code System.exit()}. In this case you have to ensure that
- * the test does not finish before {@code System.exit()} is called.
+ * the test does not finish before {@code System.exit()} is called. Use
+ * {@code ExpectedSystemExit.timeout(...)} method call to specify test method
+ * waiting time in milliseconds.
  *
  * <pre>
  * public class AppWithExit {
@@ -78,6 +80,13 @@
  *   }
  *
  *   &#064;Test
+ *   public void systemExitInSeparateThread() {
+ *     exit.expectSystemExitWithStatus(1);
+ *     exit.timeout(1000);
+ *     AppWithExit.doSomethingInSeparateThreadAndExit();
+ *   }
+ *
+ *   &#064;Test
  *   public void noSystemExit() {
  *     AppWithExit.doNothing();
  *     //passes
@@ -93,6 +102,7 @@ public static ExpectedSystemExit none() {
 	private final Collection<Assertion> assertions = new ArrayList<Assertion>();
 	private boolean expectExit = false;
 	private Integer expectedStatus = null;
+	private long timeout = 0;
 
 	private ExpectedSystemExit() {
 	}
@@ -106,6 +116,10 @@ public void expectSystemExit() {
 		expectExit = true;
 	}
 
+	public void timeout(long timeout) {
+		this.timeout = timeout;
+	}
+
 	public void checkAssertionAfterwards(Assertion assertion) {
 		assertions.add(assertion);
 	}
@@ -117,8 +131,7 @@ public Statement apply(final Statement base, Description description) {
 	}
 
 	private ProvideSecurityManager createNoExitSecurityManagerRule() {
-		NoExitSecurityManager noExitSecurityManager = new NoExitSecurityManager(
-			getSecurityManager());
+		SecurityManager noExitSecurityManager = new NoExitSecurityManager(getSecurityManager());
 		return new ProvideSecurityManager(noExitSecurityManager);
 	}
 
@@ -136,10 +149,11 @@ public void evaluate() throws Throwable {
 		};
 	}
 
-	private void checkSystemExit() {
+	private void checkSystemExit() throws InterruptedException {
 		NoExitSecurityManager securityManager = (NoExitSecurityManager) getSecurityManager();
-		if (securityManager.isCheckExitCalled())
-			handleSystemExitWithStatus(securityManager.getStatusOfFirstCheckExitCall());
+		Integer exitStatus = securityManager.getStatusOfFirstCheckExitCall(timeout);
+		if (exitStatus != null)
+			handleSystemExitWithStatus(exitStatus);
 		else
 			handleMissingSystemExit();
 	}

src/main/java/org/junit/contrib/java/lang/system/internal/NoExitSecurityManager.java

@@ -3,37 +3,32 @@
 import java.io.FileDescriptor;
 import java.net.InetAddress;
 import java.security.Permission;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.TimeUnit;
 
 /**
  * A {@code NoExitSecurityManager} throws a {@link CheckExitCalled} exception
  * whenever {@link #checkExit(int)} is called. All other method calls are
  * delegated to the original security manager.
  */
 public class NoExitSecurityManager extends SecurityManager {
+
 	private final SecurityManager originalSecurityManager;
-	private Integer statusOfFirstExitCall = null;
+	private final BlockingQueue<Integer> exitStatusHolder = new LinkedBlockingQueue<Integer>(1);
 
 	public NoExitSecurityManager(SecurityManager originalSecurityManager) {
 		this.originalSecurityManager = originalSecurityManager;
 	}
 
 	@Override
 	public void checkExit(int status) {
-		if (statusOfFirstExitCall == null)
-			statusOfFirstExitCall = status;
+		exitStatusHolder.offer(status);
 		throw new CheckExitCalled(status);
 	}
 
-	public boolean isCheckExitCalled() {
-		return statusOfFirstExitCall != null;
-	}
-
-	public int getStatusOfFirstCheckExitCall() {
-		if (isCheckExitCalled())
-			return statusOfFirstExitCall;
-		else
-			throw new IllegalStateException(
-				"checkExit(int) has not been called.");
+	public Integer getStatusOfFirstCheckExitCall(long timeout) throws InterruptedException {
+		return exitStatusHolder.poll(timeout, TimeUnit.MILLISECONDS);
 	}
 
 	@Override
@@ -44,7 +39,8 @@ public boolean getInCheck() {
 
 	@Override
 	public Object getSecurityContext() {
-		return (originalSecurityManager == null) ? super.getSecurityContext()
+		return originalSecurityManager == null
+			? super.getSecurityContext()
 			: originalSecurityManager.getSecurityContext();
 	}
 
@@ -176,7 +172,8 @@ public void checkPropertyAccess(String key) {
 
 	@Override
 	public boolean checkTopLevelWindow(Object window) {
-		return (originalSecurityManager == null) ? super.checkTopLevelWindow(window)
+		return originalSecurityManager == null
+			? super.checkTopLevelWindow(window)
 			: originalSecurityManager.checkTopLevelWindow(window);
 	}
 
@@ -230,7 +227,8 @@ public void checkSecurityAccess(String target) {
 
 	@Override
 	public ThreadGroup getThreadGroup() {
-		return (originalSecurityManager == null) ? super.getThreadGroup()
+		return originalSecurityManager == null
+			? super.getThreadGroup()
 			: originalSecurityManager.getThreadGroup();
 	}
 }

src/test/java/org/junit/contrib/java/lang/system/ExpectedSystemExitTest.java

@@ -1,21 +1,19 @@
 package org.junit.contrib.java.lang.system;
 
+import org.junit.Test;
+import org.junit.runners.model.Statement;
+
+import java.security.Permission;
+
 import static java.lang.System.getSecurityManager;
 import static java.lang.System.setSecurityManager;
 import static java.lang.Thread.sleep;
 import static org.assertj.core.api.Assertions.assertThat;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.fail;
+import static org.junit.Assert.*;
 import static org.junit.contrib.java.lang.system.Executor.exceptionThrownWhenTestIsExecutedWithRule;
 import static org.junit.contrib.java.lang.system.Executor.executeTestWithRule;
 import static org.junit.contrib.java.lang.system.Statements.TEST_THAT_DOES_NOTHING;
 
-import java.security.Permission;
-
-import org.junit.Test;
-import org.junit.runners.model.Statement;
-
 
 public class ExpectedSystemExitTest {
 	private static final Object ARBITRARY_CONTEXT = new Object();
@@ -128,7 +126,8 @@ public void current_security_manager_is_used_for_anything_else_than_system_exit(
 	@Test
 	public void test_is_successful_if_expected_exit_is_called_in_a_thread() {
 		rule.expectSystemExitWithStatus(ARBITRARY_EXIT_STATUS);
-		executeTestWithRule(new SystemExitInThread(), rule);
+		rule.timeout(1000);
+		executeTestWithRule(new SystemExitInSeparateThread(), rule);
 	}
 
 	private static class SystemExit0 extends Statement {
@@ -145,17 +144,21 @@ public void evaluate() throws Throwable {
 		}
 	}
 
-	private static class SystemExitInThread extends Statement {
+	private static class SystemExitInSeparateThread extends Statement {
 		@Override
 		public void evaluate() throws Throwable {
-			Runnable callSystemExit = new Runnable() {
-				public void run() {
+			new Thread(new LongExecutionBeforeExitCall()).start();
+		}
+
+		private static class LongExecutionBeforeExitCall implements Runnable {
+			public void run() {
+				try {
+					sleep(500);
 					System.exit(ARBITRARY_EXIT_STATUS);
+				} catch (InterruptedException e) {
+					e.printStackTrace();
 				}
-			};
-			Thread thread = new Thread(callSystemExit);
-			thread.start();
-			sleep(1000); // wait until the thread exits
+			}
 		}
 	}
 

src/test/java/org/junit/contrib/java/lang/system/internal/NoExitSecurityManagerTest.java

@@ -22,8 +22,7 @@ public class NoExitSecurityManagerTest {
 	public final TemporaryFolder temporaryFolder = new TemporaryFolder();
 
 	private final SecurityManager originalSecurityManager = mock(SecurityManager.class);
-	private final NoExitSecurityManager managerWithOriginal = new NoExitSecurityManager(
-		originalSecurityManager);
+	private final NoExitSecurityManager managerWithOriginal = new NoExitSecurityManager(originalSecurityManager);
 	private final NoExitSecurityManager managerWithoutOriginal = new NoExitSecurityManager(null);
 
 	@Test
@@ -440,21 +439,10 @@ public void getThreadGroup_may_be_called_without_original_security_manager() {
 	}
 
 	@Test
-	public void information_about_a_missing_checkExit_call_is_available() {
-		assertThat(managerWithOriginal.isCheckExitCalled()).isFalse();
-	}
-
-	@Test
-	public void information_about_a_checkExit_call_is_available() {
-		safeCallCheckExitWithStatus(DUMMY_STATUS);
-		assertThat(managerWithOriginal.isCheckExitCalled()).isTrue();
-	}
-
-	@Test
-	public void status_of_first_call_of_checkExit_is_available() {
+	public void status_of_first_call_of_checkExit_is_available() throws InterruptedException {
 		safeCallCheckExitWithStatus(DUMMY_STATUS);
 		safeCallCheckExitWithStatus(DUMMY_STATUS + 1);
-		assertThat(managerWithOriginal.getStatusOfFirstCheckExitCall())
+		assertThat(managerWithOriginal.getStatusOfFirstCheckExitCall(0))
 			.isEqualTo(DUMMY_STATUS);
 	}
 
@@ -466,14 +454,7 @@ private void safeCallCheckExitWithStatus(int status) {
 	}
 
 	@Test
-	public void fails_to_provide_status_of_first_checkExit_call_if_this_call_did_not_happen() {
-		Throwable exception = exceptionThrownBy(new Statement() {
-			public void evaluate() throws Throwable {
-				managerWithOriginal.getStatusOfFirstCheckExitCall();
-			}
-		});
-		assertThat(exception)
-			.isInstanceOf(IllegalStateException.class)
-			.hasMessage("checkExit(int) has not been called.");
+	public void null_status_of_first_checkExit_call_if_this_call_did_not_happen() throws InterruptedException {
+		assertThat(managerWithOriginal.getStatusOfFirstCheckExitCall(0)).isNull();
 	}
 }