ExpectedSystemExit improvement in multi-threads environment

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

@@ -1,18 +1,18 @@
 package org.junit.contrib.java.lang.system;
 
-import static java.lang.System.getSecurityManager;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.fail;
-
-import java.util.ArrayList;
-import java.util.Collection;
-
 import org.junit.contrib.java.lang.system.internal.CheckExitCalled;
 import org.junit.contrib.java.lang.system.internal.NoExitSecurityManager;
 import org.junit.rules.TestRule;
 import org.junit.runner.Description;
 import org.junit.runners.model.Statement;
 
+import java.util.ArrayList;
+import java.util.Collection;
+
+import static java.lang.System.getSecurityManager;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
+
 /**
  * The {@code ExpectedSystemExit} allows in-test specification of expected
  * {@code System.exit(...)} calls.
@@ -93,6 +93,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 +107,10 @@ public void expectSystemExit() {
 		expectExit = true;
 	}
 
+	public void timeout(long timeout) {
+		this.timeout = timeout;
+	}
+
 	public void checkAssertionAfterwards(Assertion assertion) {
 		assertions.add(assertion);
 	}
@@ -117,8 +122,7 @@ public Statement apply(final Statement base, Description description) {
 	}
 
 	private ProvideSecurityManager createNoExitSecurityManagerRule() {
-		NoExitSecurityManager noExitSecurityManager = new NoExitSecurityManager(
-			getSecurityManager());
+		SecurityManager noExitSecurityManager = new NoExitSecurityManager(getSecurityManager(), timeout);
 		return new ProvideSecurityManager(noExitSecurityManager);
 	}
 
@@ -138,10 +142,14 @@ public void evaluate() throws Throwable {
 
 	private void checkSystemExit() {
 		NoExitSecurityManager securityManager = (NoExitSecurityManager) getSecurityManager();
-		if (securityManager.isCheckExitCalled())
-			handleSystemExitWithStatus(securityManager.getStatusOfFirstCheckExitCall());
-		else
-			handleMissingSystemExit();
+		try {
+			if (securityManager.isCheckExitCalled())
+                handleSystemExitWithStatus(securityManager.getStatusOfFirstCheckExitCall());
+            else
+                handleMissingSystemExit();
+		} catch (InterruptedException e) {
+			fail("Unexpected waiting loop break.");
+		}
 	}
 
 	private void handleMissingSystemExit() {

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

@@ -3,37 +3,45 @@
 import java.io.FileDescriptor;
 import java.net.InetAddress;
 import java.security.Permission;
+import java.util.concurrent.CountDownLatch;
+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 final long timeout;
+
+	private final CountDownLatch synchLatch = new CountDownLatch(1);
 	private Integer statusOfFirstExitCall = null;
 
-	public NoExitSecurityManager(SecurityManager originalSecurityManager) {
+	public NoExitSecurityManager(SecurityManager originalSecurityManager, long timeout) {
 		this.originalSecurityManager = originalSecurityManager;
+		this.timeout = timeout;
 	}
 
 	@Override
 	public void checkExit(int status) {
 		if (statusOfFirstExitCall == null)
 			statusOfFirstExitCall = status;
+		synchLatch.countDown();
 		throw new CheckExitCalled(status);
 	}
 
-	public boolean isCheckExitCalled() {
-		return statusOfFirstExitCall != null;
+	public boolean isCheckExitCalled() throws InterruptedException {
+		boolean wasCountDown = synchLatch.await(timeout, TimeUnit.MILLISECONDS);
+		return statusOfFirstExitCall != null && wasCountDown;
 	}
 
-	public int getStatusOfFirstCheckExitCall() {
+	public int getStatusOfFirstCheckExitCall() throws InterruptedException {
 		if (isCheckExitCalled())
 			return statusOfFirstExitCall;
 		else
-			throw new IllegalStateException(
-				"checkExit(int) has not been called.");
+			throw new IllegalStateException("checkExit(int) has not been called.");
 	}
 
 	@Override

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,9 +22,8 @@ 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 managerWithoutOriginal = new NoExitSecurityManager(null);
+	private final NoExitSecurityManager managerWithOriginal = new NoExitSecurityManager(originalSecurityManager, 0);
+	private final NoExitSecurityManager managerWithoutOriginal = new NoExitSecurityManager(null, 0);
 
 	@Test
 	public void an_exception_with_the_status_is_thrown_when_checkExit_is_called() {
@@ -440,18 +439,18 @@ public void getThreadGroup_may_be_called_without_original_security_manager() {
 	}
 
 	@Test
-	public void information_about_a_missing_checkExit_call_is_available() {
+	public void information_about_a_missing_checkExit_call_is_available() throws InterruptedException {
 		assertThat(managerWithOriginal.isCheckExitCalled()).isFalse();
 	}
 
 	@Test
-	public void information_about_a_checkExit_call_is_available() {
+	public void information_about_a_checkExit_call_is_available() throws InterruptedException {
 		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())