Skip to Content

Testing System.exit() with JUnit5

An Extension to JUnit 5 for testing System.exit(), and how it works

Posted on by Todd Ginsberg
Photo by Ricardo Gomez Angel on Unsplash
Photo by Ricardo Gomez Angel on Unsplash

Testing System.exit()

In all my years of writing Java, I haven’t called System.exit() in production code too many times. However, every once in a while I will encounter a Perfectly Legitimate Reason™ to do so. Once I write the code, the next most obvious question is - “How am I going to actually test this?” In this post, I’ll explore how to safely test code that calls System.exit() with JUnit 5.

The code for this post, as well as instructions for using it in your own projects are hosted in this GitHub repository, please check it out!

Update (2021-02-21): Version 1.1.1 is out on Maven Central, or you can check it out on GitHub.

Preventing an Actual System Exit

Writing a unit test that actually exits the JVM while it is under test is definitely not ideal. Whenever the JVM does something interesting (like exiting, or reading a file), it first checks whether it has permission to do so. This is done by consulting the SecurityManager the system is using. One of the methods on SecurityManager is checkExit(). If that method throws a SecurityException, it means the system is not allowed to exit at this point. Conveniently, checkExit() takes one argument - the exit status code being attempted.

Now that we know that, we can form a plan of action:

  1. Write a SecurityManager that always prevents System.exit(), and records the code attempted.
  2. Any other call to our SecurityManager should delegate to whatever SecurityManager was being used by the system before our test started.
  3. Integrate this with JUnit 5, via the Extension model.

Custom SecurityManager

Writing a SecurityManager that prevents System.exit() and records the code that was used is pretty straight forward now that we know how it works:

public class DisallowExitSecurityManager extends SecurityManager {
    private final SecurityManager delegatedSecurityManager;
    private Integer firstExitStatusCode;

    public DisallowExitSecurityManager(final SecurityManager originalSecurityManager) {
        this.delegatedSecurityManager = originalSecurityManager;
    }

    /**
     * This is the one method we truly override in this class, all others are delegated.
     *
     * @param statusCode the exit status
     */
    @Override
    public void checkExit(final int statusCode) {
        if (firstExitStatusCode == null) {
            this.firstExitStatusCode = statusCode;
        }
        throw new SystemExitPreventedException();
    }

    public Integer getFirstExitStatusCode() {
        return firstExitStatusCode;
    }

    // All other methods implemented and delegate to delegatedSecurityManager, if it is present. 
    // Otherwise, they do nothing and allow the check to pass.

    // Example:
    @Override
    public void checkPermission(Permission perm) {
        if (delegatedSecurityManager != null) {
            delegatedSecurityManager.checkPermission(perm);
        }
    }
}

As you can see, when the system checks permission to exit, we save the code being used in firstExitStatusCode, because it only makes sense to only record the first attempt. Then we throw a SystemExitPreventedException, which we have defined for this purpose. We define our own exception because we need to detect this exact scenario and don’t want to inadvertently handle a more general SecurityException which might mean something else entirely. In short, we define our own exception here because we can identify it later.

For the sake of completeness, here is our SystemExitPreventedException, which is a simple subclass of SecurityException. In theory we could record the security code here use it when we catch this exception, but I didn’t want to make this exception stateful.

class SystemExitPreventedException extends SecurityException {

}

The JUnit 5 Extension

One of the changes in JUnit 5 is the introduction of the Extension Model . This one extension point takes over for @Rule, @ClassRule, and Runners. Extension itself is just a marker interface, and there are several ways to use it.

In order to register our Extension, we will create two annotations that we will use in our test cases: one for asserting that some code should call System.exit() with any code, and one that specifies a code. It is quite convenient that JUnit5 supports meta-annotations (annotations on annotations) so we can reference the actual working part of our code:

@Retention(RetentionPolicy.RUNTIME)
@Target({ElementType.TYPE, ElementType.METHOD, ElementType.ANNOTATION_TYPE})
@ExtendWith(SystemExitExtension.class)
public @interface ExpectSystemExit {

}

@Retention(RetentionPolicy.RUNTIME)
@Target({ElementType.TYPE, ElementType.METHOD, ElementType.ANNOTATION_TYPE})
@ExtendWith(SystemExitExtension.class)
public @interface ExpectSystemExitWithStatus {
    int value();
}

As you can see, each of our annotations are themselves annotated with @ExtendWith(SystemExitExtension.class). This tells the Jupiter runtime (JUnit 5’s name for itself) to look for a class called SystemExitExtension and register it.

Before we create this extension, let’s decide what it needs to do:

  1. Before each test: replace the existing SecurityManager (if any) with the one we wrote above.
  2. During each test: catch and swallow our custom exception (SystemExitPreventedException) and re-throw all others.
  3. After each test: return the original SecurityManager (if there was one) to service, and remove our custom one.

Because the Extension model is so granular, we will need to implement three interfaces with one method each: BeforeEachCallback, TestExecutionExceptionHandler, and AfterEachCallback.

We will also need a method to find the ExpectSystemExitWithStatus annotation (otherwise we can just assume that our extension was called as a result of ExpectedsystemExit and we don’t need to worry about a specific code).

public class SystemExitExtension implements BeforeEachCallback, AfterEachCallback, TestExecutionExceptionHandler {
    private Integer expectedStatusCode;  
    private final DisallowExitSecurityManager disallowExitSecurityManager = new DisallowExitSecurityManager(System.getSecurityManager());
    private SecurityManager originalSecurityManager;

    @Override
    public void beforeEach(final ExtensionContext context) {
        // TODO
    }

    @Override
    public void handleTestExecutionException(final ExtensionContext context, final Throwable throwable) 
        throws Throwable {
        // TODO
    }

    @Override
    public void afterEach(final ExtensionContext context) {
        // TODO
    }

    // Find the annotation on a method, or failing that, a class.
    private Optional<ExpectSystemExitWithStatus> getAnnotation(final ExtensionContext context) {
        final Optional<ExpectSystemExitWithStatus> method = 
            findAnnotation(context.getTestMethod(), ExpectSystemExitWithStatus.class);

        if (method.isPresent()) {
            return method;
        } else {
            return findAnnotation(context.getTestClass(), ExpectSystemExitWithStatus.class);
        }
    }
}

The strategy we will follow for finding our annotation is to search the test method we are executing for it, and failing that, the test class itself. This gives us some flexibility to annotate a specific test or a group of tests (even nested) that a System.exit() with a specific code should be expected. In this method, we use findAnnotation, which is included in and used by JUnit 5.

In the code above, we’ve already defined all of the instance variables we will need so we can move on to implementing our three lifecycle methods (before, during, after).

Before Test Execution: BeforeEachCallback

@Override
public void beforeEach(final ExtensionContext context) {
    // Set aside the current SecurityManager
    originalSecurityManager = System.getSecurityManager();

    // Get the expected exit status code, if any
    getAnnotation(context).ifPresent(code -> expectedStatusCode = code.value());

    // Install our own SecurityManager
    System.setSecurityManager(disallowExitSecurityManager);
}

As you can see, we set aside the currently running system SecurityManager, even if it is null. Then we find our annotation and capture the expectedStatusCode if there is one. Finally, we tell Java to use the DisallowExitSecurityManager that we wrote for the duration of the test.

During Test Execution: TestExecutionExceptionHandler

public void handleTestExecutionException(final ExtensionContext context, final Throwable throwable) 
    throws Throwable {

    if (!(throwable instanceof SystemExitPreventedException)) {
        throw throwable;
    }
}

In this implementation we will swallow our custom exception (SystemExitPreventedException)) and re-throw anything else because those would be legitimate exceptions thrown by the original SecurityManager. Why are we throwing an exception only to catch and ignore it? Because the JVM needs us to throw some kind of SecurityException out of checkExit() or it will actually exit (which would be bad). And since JUnit will fail on any exception, we need to swallow it and prevent it from bubbling up our call stack to end our test.

After Test Execution: SystemExitPreventedException

@Override
public void afterEach(final ExtensionContext context) {
    // Return the original SecurityManager, if any, to service.
    System.setSecurityManager(originalSecurityManager);

    if (expectedStatusCode == null) {
        assertNotNull(
                disallowExitSecurityManager.getFirstExitStatusCode(),
                "Expected System.exit() to be called, but it was not"
        );
    } else {
        assertEquals(
                expectedStatusCode,
                disallowExitSecurityManager.getFirstExitStatusCode(),
                "Expected System.exit(" + expectedStatusCode + ") to be called, but it was not."
        );
    }
}

This one is a bit more complicated because we have to determine if we’re expecting any exit code or a specific exit code. First, we have to return the original SecurityManager to service, replacing our custom implementation. Next, we use JUnit’s own assertions to test our condition - that System.exit() was called, possibly with a specific code. We can get both of these pieces of information from our SecurityManager, which we still hold a reference do.

Thoughts

I really like how the Extension Model breaks lifecycle events down into simple methods to be implemented or not. This model seems a lot more coherent than the concepts it is designed to replace. Yes, in the short term we won’t be able to use any @Rules that may already exit from JUnit 4, but I feel this is a lot more expressive. I’m eager to see what the community develops in order to assist with testing as we move to JUnit 5!

Using The Extension

Now that our extension is written (and tested!) we can write our own tests that consume it:

public class MyTestCases { 
    
    @Test
    @ExpectSystemExit
    public void thatSystemExitIsCalled() {
        System.exit(1);
    }

    @Test
    @ExpectSystemExitWithStatus(42)
    public void thatSystemExitIsCalledWithSpecificCode() {
        System.exit(42);
    }
}

We could also annotate our test classes themselves, rather than the methods individually if we had identical conditions.

Using junit5-system-exit

As I mentioned in the into to this post, I’ve already written this code and published it as a dependency you can get from Maven Central to use in your projects. To use it in your build, just add one of these to your pom.xml or build.gradle:

// Gradle build.gradle
testImplementation("com.ginsberg:junit5-system-exit:1.1.1")

<!-- Maven pom.xml -->
<dependency>
    <groupId>com.ginsberg</groupId>
    <artifactId>junit5-system-exit</artifactId>
    <version>1.1.1</version>
    <scope>test</scope>
</dependency>

I hope it helps somebody else - let me know if you use this in your project. I’ve released this under the MIT License , so it should be safe for corporate use, generally speaking.