Mocking ClassLoader in Java Tests Without Breaking the Real ClassLoader

Question

I encountered an issue with mocking the ClassLoader in my tests. The method under test uses Thread.currentThread().getContextClassLoader() to load a resource, and when I mock the ClassLoader, it breaks the real ClassLoader, causing subsequent tests to fail. I'm looking for a way to mock or use the ClassLoader without breaking the real one. Specifically, I want to either obtain the real ClassLoader before the test and set it back after the test, or find a way to create a file and use it in the test without affecting the real ClassLoader. How can I refactor my test to achieve this?

The method from utility class

public static KeyPair getKeyPairFromKeystore(JwtConfigProperties jwtConfigProperties,
                                                 KeystoreConfigProperties keystoreConfigProperties) {
        log.info( "Called get keypair from keystore");

        KeyStore.PrivateKeyEntry privateKeyEntry = null;
        Certificate cert = null;

        String keystoreName = keystoreConfigProperties.getName();

        try (InputStream in = Thread.currentThread()
                .getContextClassLoader()
                .getResourceAsStream("security/" + keystoreName)
        ) {
            if (in == null) {
                log.info("Input file is null!");
                throw new NoSuchElementException("Input file  is null");
            }

            char[] keystorePassword = keystoreConfigProperties.getPassword().toCharArray();
            char[] jwtPassword = jwtConfigProperties.getPassword().toCharArray();
            String jwtAlias = jwtConfigProperties.getAlias();

            //enforcing jceks
            KeyStore keyStore = KeyStore.getInstance("JCEKS");
            keyStore.load(in, keystorePassword);
            KeyStore.PasswordProtection keyPassword = new KeyStore
                    .PasswordProtection(jwtPassword);
            privateKeyEntry = (KeyStore.PrivateKeyEntry) keyStore.getEntry(jwtAlias, keyPassword);
            cert = keyStore.getCertificate(jwtAlias);
            log.debug("Public key: {}", cert.getPublicKey().toString());
        } catch (KeyStoreException | IOException | NoSuchAlgorithmException
                | CertificateException | UnrecoverableEntryException e) {
            log.error("Error message: {}, Exception: {}", e.getMessage(), e);
        }

        return new KeyPair(Objects.requireNonNull(cert).getPublicKey(),
            Objects.requireNonNull(privateKeyEntry).getPrivateKey());
    }

JwtConfigProperties & KeystoreConfigProperties it`s configuration properties from application.yml

My test

@RunWith(MockitoJUnitRunner.class)
public class KeystoreUtilTest {


    @Mock
    private ClassLoader newLoader;

    @Mock
    private InputStream in;

    @Mock
    private KeyStore keystore;


    @Test
    public void testGetKeyPairFromKeystore_validPublicKey() throws KeyStoreException {

        Certificate certificate = mock(Certificate.class);
        PublicKey publicKey = mock(PublicKey.class);

        JwtConfigProperties jwtConfigProperties = new JwtConfigProperties();
        jwtConfigProperties.setAlias("alias");
        jwtConfigProperties.setPassword("password");
        KeystoreConfigProperties keystoreConfigProperties = new KeystoreConfigProperties();
        keystoreConfigProperties.setName("name");
        keystoreConfigProperties.setPassword("password");
        Thread.currentThread().setContextClassLoader(newLoader);

        when(newLoader.getResourceAsStream("security/" + keystoreConfigProperties.getName())).thenReturn(in);

        try (MockedStatic<KeyStore> keyStoreMockedStatic = mockStatic(KeyStore.class)) {

            keyStoreMockedStatic
                    .when(() -> KeyStore.getInstance("JCEKS"))
                    .thenReturn(keystore);

            keyStoreMockedStatic
                    .when(() -> KeyStore.getInstance("JCEKS").getCertificate("alias"))
                    .thenReturn(certificate);

            keyStoreMockedStatic
                    .when(() -> KeyStore.getInstance("JCEKS").getCertificate("alias").getPublicKey())
                    .thenReturn(publicKey);

            assertThrows(NullPointerException.class,
                    () -> KeystoreUtil.getKeyPairFromKeystore(jwtConfigProperties, keystoreConfigProperties));
            verify(keystore, times(2)).getCertificate("alias");
            keyStoreMockedStatic
                    .verify(() -> KeyStore.getInstance("JCEKS"), times(3));
        }

    }

In the test, I am mocking Classloader, but the real one stopped his work and tests don't start after this test

Answer 1

Don't mock it. Inject a function as a parameter.

public static KeyPair getKeyPairFromKeystore(JwtConfigProperties jwtConfigProperties,
                                             KeystoreConfigProperties keystoreConfigProperties,
                                             Function<String, InputStream> getInputStream
) {

    //...
    try (InputStream in = getInputStream.apply("security/" + keystoreName)) {
        //...
    }
}

In your real code, pass a function

getKeyPairFromKeystore(..., ...,
    str -> Thread.currentThread()
        .getContextClassLoader()
        .getResourceAsStream(str)
)

In your test, pass whatever function returns the InputStream you want for the test.