Refactors and tests in java-spiffe-helper.

Signed-off-by: Max Lambrecht <maxlambrecht@gmail.com>

java-spiffe-core/src/main/java/io/spiffe/workloadapi/JwtSource.java

@@ -168,8 +168,6 @@ public class JwtSource implements JwtSvidSource, BundleSource<JwtBundle>, Closea
         }
     }
 
-
-
     private void init(final Duration timeout) throws TimeoutException {
         CountDownLatch done = new CountDownLatch(1);
         setJwtBundlesWatcher(done);

java-spiffe-helper/src/main/java/io/spiffe/helper/cli/Runner.java

@@ -1,6 +1,8 @@
 package io.spiffe.helper.cli;
 
 import io.spiffe.exception.SocketEndpointAddressException;
+import io.spiffe.helper.exception.KeyStoreHelperException;
+import io.spiffe.helper.exception.RunnerException;
 import io.spiffe.helper.keystore.KeyStoreHelper;
 import io.spiffe.helper.keystore.KeyStoreType;
 import lombok.extern.java.Log;
@@ -18,7 +20,6 @@ import java.io.InputStream;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
-import java.security.KeyStoreException;
 import java.util.Properties;
 
 /**
@@ -37,41 +38,45 @@ public class Runner {
      *
      * @param args contains the option with the config file path
      */
-    public static void main(final String ...args) {
+    public static void main(final String ...args) throws RunnerException {
         String configFilePath = null;
         try {
             configFilePath = getCliConfigOption(args);
             val parameters = parseConfigFile(Paths.get(configFilePath));
             val options = toKeyStoreOptions(parameters);
-
+            try (val keyStoreHelper = KeyStoreHelper.create(options)) {
-            // run KeyStoreHelper
+                keyStoreHelper.run(true);
-            new KeyStoreHelper(options);
+            }
-        } catch (IOException e) {
+        } catch (SocketEndpointAddressException | KeyStoreHelperException | RunnerException | IllegalArgumentException e) {
-            log.severe(String.format("Cannot open config file: %s %n %s", configFilePath, e.getMessage()));
+            log.severe(e.getMessage());
-        } catch (KeyStoreException e) {
+            throw new RunnerException(e);
-            log.severe(String.format("Error storing certs in keystores: %s", e.getMessage()));
-        } catch (SocketEndpointAddressException e) {
-            log.severe(String.format("Workload API address is not valid: %s", e.getMessage()));
-        } catch (ParseException e) {
-            log.severe(String.format( "%s. Use -c, --config <arg>", e.getMessage()));
         }
     }
 
-    static Properties parseConfigFile(final Path configFilePath) throws IOException {
+    static Properties parseConfigFile(final Path configFilePath) throws RunnerException {
         Properties prop = new Properties();
         try (InputStream in = Files.newInputStream(configFilePath)) {
             prop.load(in);
+        } catch (IOException e) {
+            val error = String.format("Cannot open config file: %s %n %s", configFilePath, e.getMessage());
+            throw new RunnerException(error);
         }
         return prop;
     }
 
-    static String getCliConfigOption(final String ...args) throws ParseException {
+    static String getCliConfigOption(final String ...args) throws RunnerException {
         final Options cliOptions = new Options();
         final Option confOption = new Option("c", "config", true, "config file");
         confOption.setRequired(true);
         cliOptions.addOption(confOption);
         CommandLineParser parser = new DefaultParser();
-        CommandLine cmd = parser.parse(cliOptions, args);
+        CommandLine cmd = null;
+        try {
+            cmd = parser.parse(cliOptions, args);
+        } catch (ParseException e) {
+            val error = String.format( "%s. Use -c, --config <arg>", e.getMessage());
+            throw new RunnerException(error);
+        }
         return cmd.getOptionValue("config");
     }
 

java-spiffe-helper/src/main/java/io/spiffe/helper/exception/KeyStoreHelperException.java

@@ -0,0 +1,15 @@
+package io.spiffe.helper.exception;
+
+/**
+ * Checked exception to be thrown when there is an error creating or initializing a KeyStoreHelper.
+ */
+public class KeyStoreHelperException extends Exception {
+
+    public KeyStoreHelperException(String message) {
+        super(message);
+    }
+
+    public KeyStoreHelperException(String message, Throwable cause) {
+        super(message, cause);
+    }
+}

java-spiffe-helper/src/main/java/io/spiffe/helper/exception/RunnerException.java

@@ -0,0 +1,14 @@
+package io.spiffe.helper.exception;
+
+/**
+ * Checked exception to be thrown when there are errors configuring the cli Runner.
+ */
+public class RunnerException extends Exception {
+    public RunnerException(String message) {
+        super(message);
+    }
+
+    public RunnerException(Throwable cause) {
+        super(cause);
+    }
+}

java-spiffe-helper/src/main/java/io/spiffe/helper/keystore/KeyStoreHelper.java

@@ -2,6 +2,8 @@ package io.spiffe.helper.keystore;
 
 import io.spiffe.bundle.x509bundle.X509Bundle;
 import io.spiffe.exception.SocketEndpointAddressException;
+import io.spiffe.exception.WatcherException;
+import io.spiffe.helper.exception.KeyStoreHelperException;
 import io.spiffe.spiffeid.TrustDomain;
 import io.spiffe.workloadapi.DefaultWorkloadApiClient;
 import io.spiffe.workloadapi.Watcher;
@@ -53,90 +55,131 @@ public class KeyStoreHelper implements Closeable {
 
     private final WorkloadApiClient workloadApiClient;
 
+    private volatile boolean closed;
+    private volatile CountDownLatch countDownLatch;
 
     /**
-     * Constructor.
-     * <p>
      * Creates an instance of a KeyStoreHelper for fetching X.509 SVIDs and bundles
      * from a Workload API and store them in a binary Java KeyStore in disk.
-     * <p>
-     * It blocks until the initial update has been received from the Workload API.
      *
      * @param options an instance of {@link KeyStoreOptions}
-     * @throws SocketEndpointAddressException is the socket endpoint address is not valid
+     * @return an instance of a KeyStoreHelper
-     * @throws KeyStoreException              is the entry cannot be stored in the KeyStore
+     * @throws SocketEndpointAddressException if the socket endpoint address is not valid
+     * @throws KeyStoreHelperException        if the KeyStoreHelper cannot be created
      */
-    public KeyStoreHelper(@NonNull final KeyStoreOptions options)
+    public static KeyStoreHelper create(@NonNull final KeyStoreOptions options) throws SocketEndpointAddressException, KeyStoreHelperException {
-            throws SocketEndpointAddressException, KeyStoreException {
-
-        final KeyStoreType keyStoreType;
-        if (options.keyStoreType == null) {
-            keyStoreType = KeyStoreType.getDefaultType();
-        } else {
-            keyStoreType = options.keyStoreType;
-        }
-
-        this.keyPass = options.keyPass;
-
-        if (StringUtils.isBlank(options.keyAlias)) {
-            this.keyAlias = DEFAULT_ALIAS;
-        } else {
-            this.keyAlias = options.keyAlias;
-        }
 
         if (options.keyStorePath.equals(options.trustStorePath)) {
-            throw new KeyStoreException("KeyStore and TrustStore should use different files");
+            throw new KeyStoreHelperException("KeyStore and TrustStore should use different files");
         }
 
-        this.keyStore = KeyStore.builder()
+        if (options.keyStoreType == null) {
-                .keyStoreFilePath(options.keyStorePath)
+            options.keyStoreType = KeyStoreType.getDefaultType();
-                .keyStoreType(keyStoreType)
-                .keyStorePassword(options.keyStorePass)
-                .build();
-
-        this.trustStore = KeyStore.builder()
-                .keyStoreFilePath(options.trustStorePath)
-                .keyStoreType(keyStoreType)
-                .keyStorePassword(options.trustStorePass)
-                .build();
-
-        if (options.workloadApiClient != null) {
-            workloadApiClient = options.workloadApiClient;
-        } else {
-            workloadApiClient = createNewClient(options.spiffeSocketPath);
         }
 
-        setX509ContextWatcher(workloadApiClient);
+        if (StringUtils.isBlank(options.keyAlias)) {
+            options.keyAlias = DEFAULT_ALIAS;
+        }
+
+        val keyStore = createKeyStore(options, options.keyStorePath, options.keyStorePass);
+        val trustStore = createKeyStore(options, options.trustStorePath, options.trustStorePass);
+
+        if (options.workloadApiClient == null) {
+            options.workloadApiClient = createNewClient(options.spiffeSocketPath);
+        }
+
+        return new KeyStoreHelper(keyStore, trustStore, options.keyPass, options.keyAlias, options.workloadApiClient);
     }
 
+    /**
+     * Sets the instance to run fetching and storing the X.509 SVIDs and Bundles.
+     *
+     * @param keepRunning if true, the process will block receiving and storing updates, otherwise it blocks only until
+     *                    the first X.509 context is received and stored.
+     * @throws KeyStoreHelperException if there is an error fetching or storing the X.509 SVIDs and Bundles
+     */
+    public void run(boolean keepRunning) throws KeyStoreHelperException {
+        if (isClosed()) {
+            throw new IllegalStateException("KeyStoreHelper is closed");
+        }
+
+        try {
+            this.setX509ContextWatcher(keepRunning);
+        } catch (Exception e) {
+            throw new KeyStoreHelperException("Error running KeyStoreHelper", e);
+        }
+    }
+
+    /**
+     * Closes the KeyStoreHelper instance.
+     */
     @SneakyThrows
     @Override
     public void close() {
-        workloadApiClient.close();
+        if (!closed) {
+            synchronized (this) {
+                if (!closed) {
+                    workloadApiClient.close();
+                    countDown();
+                    closed = true;
+                    log.info("KeyStoreHelper is closed");
+                }
+            }
+        }
     }
 
+    private void countDown() {
+        if (countDownLatch != null) {
+            countDownLatch.countDown();
+        }
+    }
 
-    private WorkloadApiClient createNewClient(final String spiffeSocketPath) throws SocketEndpointAddressException {
+    private KeyStoreHelper(KeyStore keyStore, KeyStore trustStore, String keyPass, String keyAlias, WorkloadApiClient workloadApiClient) {
+        this.keyStore = keyStore;
+        this.trustStore = trustStore;
+        this.keyPass = keyPass;
+        this.keyAlias = keyAlias;
+        this.workloadApiClient = workloadApiClient;
+    }
+
+    private static KeyStore createKeyStore(KeyStoreOptions options, Path keyStorePath, String keyStorePass) throws KeyStoreHelperException {
+        try {
+            return KeyStore.builder()
+                    .keyStoreFilePath(keyStorePath)
+                    .keyStoreType(options.keyStoreType)
+                    .keyStorePassword(keyStorePass)
+                    .build();
+        } catch (KeyStoreException e) {
+            throw new KeyStoreHelperException("Error creating KeyStore/TrustStore", e);
+        }
+    }
+
+    private static WorkloadApiClient createNewClient(final String spiffeSocketPath) throws SocketEndpointAddressException {
         val clientOptions = DefaultWorkloadApiClient.ClientOptions.builder().spiffeSocketPath(spiffeSocketPath).build();
         return DefaultWorkloadApiClient.newClient(clientOptions);
     }
 
-    private void setX509ContextWatcher(final WorkloadApiClient workloadApiClient) {
+    private void setX509ContextWatcher(boolean keepRunning) {
-        val countDownLatch = new CountDownLatch(1);
+        countDownLatch = new CountDownLatch(1);
         workloadApiClient.watchX509Context(new Watcher<X509Context>() {
             @Override
             public void onUpdate(X509Context update) {
                 try {
                     storeX509ContextUpdate(update);
+                    if (!keepRunning) {
+                        // got a X509 update, process is complete
+                        countDownLatch.countDown();
+                    }
                 } catch (KeyStoreException e) {
                     this.onError(e);
                 }
-                countDownLatch.countDown();
             }
 
             @Override
-            public void onError(Throwable t) {
+            public void onError(Throwable e) {
-                log.log(Level.SEVERE, "Error processing X.509 context update", t);
+                log.log(Level.SEVERE, e.getMessage());
+                countDownLatch.countDown();
+                throw new WatcherException("Error processing X.509 context update", e);
             }
         });
 
@@ -177,6 +220,12 @@ public class KeyStoreHelper implements Closeable {
         return trustDomain.getName().concat(".").concat(String.valueOf(index));
     }
 
+    private boolean isClosed() {
+        synchronized (this) {
+            return closed;
+        }
+    }
+
     private void await(final CountDownLatch latch) {
         try {
             latch.await();
@@ -200,7 +249,7 @@ public class KeyStoreHelper implements Closeable {
      * for information about standard keystore types.
      * <p>
      * The same type is used for both the KeyStore and the TrustStore.
-     *
+     * <p>
      * Optional. Default is PKCS12.
      * <p>
      * <code>keyStorePass</code> The password to generate the keystore integrity check.

java-spiffe-helper/src/main/java/io/spiffe/helper/keystore/PrivateKeyEntry.java

@@ -1,17 +1,27 @@
 package io.spiffe.helper.keystore;
 
+import lombok.AccessLevel;
 import lombok.Builder;
-import lombok.Value;
+import lombok.Data;
+import lombok.Setter;
 
 import java.security.Key;
 import java.security.cert.X509Certificate;
 
-@Value
+@Data
 class PrivateKeyEntry {
-    String alias;
+
-    Key privateKey;
+    @Setter(AccessLevel.NONE)
-    String password;
+    private String alias;
-    X509Certificate[] certificateChain;
+
+    @Setter(AccessLevel.NONE)
+    private Key privateKey;
+
+    @Setter(AccessLevel.NONE)
+    private String password;
+
+    @Setter(AccessLevel.NONE)
+    private X509Certificate[] certificateChain;
 
     @Builder
     PrivateKeyEntry(

java-spiffe-helper/src/test/java/io/spiffe/helper/cli/RunnerTest.java

@@ -1,17 +1,16 @@
 package io.spiffe.helper.cli;
 
 import io.spiffe.exception.SocketEndpointAddressException;
-import org.apache.commons.cli.ParseException;
+import io.spiffe.helper.exception.RunnerException;
 import org.junit.jupiter.api.Test;
 
-import java.io.IOException;
-import java.net.URI;
 import java.net.URISyntaxException;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.security.KeyStoreException;
 import java.util.Properties;
 
+import static io.spiffe.helper.utils.TestUtils.toUri;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.fail;
 
@@ -23,8 +22,8 @@ class RunnerTest {
         try {
             Runner.main("-c", path.toString());
             fail("expected exception: property is missing");
-        } catch (IllegalArgumentException e) {
+        } catch (RunnerException e) {
-            assertEquals("keyStorePath config is missing", e.getMessage());
+            assertEquals("keyStorePath config is missing", e.getCause().getMessage());
         }
     }
 
@@ -34,8 +33,8 @@ class RunnerTest {
         try {
             Runner.main("-c", path.toString());
             fail("expected exception: property is missing");
-        } catch (IllegalArgumentException e) {
+        } catch (RunnerException e) {
-            assertEquals("keyStorePass config is missing", e.getMessage());
+            assertEquals("keyStorePass config is missing", e.getCause().getMessage());
         }
     }
 
@@ -45,8 +44,8 @@ class RunnerTest {
         try {
             Runner.main("-c", path.toString());
             fail("expected exception: property is missing");
-        } catch (IllegalArgumentException e) {
+        } catch (RunnerException e) {
-            assertEquals("keyPass config is missing", e.getMessage());
+            assertEquals("keyPass config is missing", e.getCause().getMessage());
         }
     }
 
@@ -56,8 +55,8 @@ class RunnerTest {
         try {
             Runner.main("-c", path.toString());
             fail("expected exception: property is missing");
-        } catch (IllegalArgumentException e) {
+        } catch (RunnerException e) {
-            assertEquals("trustStorePath config is missing", e.getMessage());
+            assertEquals("trustStorePath config is missing", e.getCause().getMessage());
         }
     }
 
@@ -67,8 +66,8 @@ class RunnerTest {
         try {
             Runner.main("-c", path.toString());
             fail("expected exception: property is missing");
-        } catch (IllegalArgumentException e) {
+        } catch (RunnerException e) {
-            assertEquals("trustStorePass config is missing", e.getMessage());
+            assertEquals("trustStorePass config is missing", e.getCause().getMessage());
         }
     }
 
@@ -77,7 +76,7 @@ class RunnerTest {
         String option = null;
         try {
             option = Runner.getCliConfigOption("-c", "example");
-        } catch (ParseException e) {
+        } catch (RunnerException e) {
             fail(e);
         }
         assertEquals("example", option);
@@ -88,7 +87,7 @@ class RunnerTest {
         String option = null;
         try {
             option = Runner.getCliConfigOption("--config", "example");
-        } catch (ParseException e) {
+        } catch (RunnerException e) {
             fail(e);
         }
         assertEquals("example", option);
@@ -99,13 +98,13 @@ class RunnerTest {
         try {
             Runner.getCliConfigOption("--unknown", "example");
             fail("expected parse exception");
-        } catch (ParseException e) {
+        } catch (RunnerException e) {
-            assertEquals("Unrecognized option: --unknown", e.getMessage());
+            assertEquals("Unrecognized option: --unknown. Use -c, --config <arg>", e.getMessage());
         }
     }
 
     @Test
-    void test_ParseConfigFile() throws URISyntaxException, IOException {
+    void test_ParseConfigFile() throws URISyntaxException, RunnerException {
         final Path path = Paths.get(toUri("testdata/cli/correct.conf"));
         final Properties properties = Runner.parseConfigFile(path);
 
@@ -118,8 +117,4 @@ class RunnerTest {
         assertEquals("other_alias", properties.getProperty("keyAlias"));
         assertEquals("unix:/tmp/agent.sock", properties.getProperty("spiffeSocketPath"));
     }
-
-    private URI toUri(String path) throws URISyntaxException {
-        return Thread.currentThread().getContextClassLoader().getResource(path).toURI();
-    }
 }

java-spiffe-helper/src/test/java/io/spiffe/helper/keystore/KeyStoreHelperTest.java

@@ -1,8 +1,11 @@
 package io.spiffe.helper.keystore;
 
 import io.spiffe.exception.SocketEndpointAddressException;
+import io.spiffe.helper.exception.KeyStoreHelperException;
+import io.spiffe.helper.utils.TestUtils;
 import io.spiffe.internal.CertificateUtils;
 import io.spiffe.spiffeid.SpiffeId;
+import io.spiffe.workloadapi.Address;
 import io.spiffe.workloadapi.WorkloadApiClient;
 import lombok.SneakyThrows;
 import lombok.val;
@@ -29,16 +32,25 @@ import static org.junit.jupiter.api.Assertions.fail;
 class KeyStoreHelperTest {
 
     private WorkloadApiClient workloadApiClient;
+    private WorkloadApiClient workloadApiClientErrorStub;
     private Path keyStoreFilePath;
     private Path trustStoreFilePath;
 
     @BeforeEach
     void setUp() {
         workloadApiClient = new WorkloadApiClientStub();
+        workloadApiClientErrorStub = new WorkloadApiClientErrorStub();
+    }
+
+    @SneakyThrows
+    @AfterEach
+    void tearDown() {
+        deleteFile(keyStoreFilePath);
+        deleteFile(trustStoreFilePath);
     }
 
     @Test
-    void testNewHelper_certs_are_stored_successfully() throws KeyStoreException, SocketEndpointAddressException, UnrecoverableKeyException, CertificateException, NoSuchAlgorithmException, IOException {
+    void testNewHelper_certs_are_stored_successfully() throws SocketEndpointAddressException, UnrecoverableKeyException, CertificateException, NoSuchAlgorithmException, IOException, KeyStoreException {
 
         val keyStorefileName = RandomStringUtils.randomAlphabetic(10);
         keyStoreFilePath = Paths.get(keyStorefileName);
@@ -65,7 +77,11 @@ class KeyStoreHelperTest {
                 .build();
 
         // run KeyStoreHelper
-        new KeyStoreHelper(options);
+        try (val keystoreHelper = KeyStoreHelper.create(options)) {
+            keystoreHelper.run(false);
+        } catch (KeyStoreHelperException e) {
+            fail();
+        }
 
         checkPrivateKeyEntry(keyStoreFilePath, keyStorePass, keyPass, keyStoreType, alias);
         val authority1Alias = "example.org.0";
@@ -95,8 +111,11 @@ class KeyStoreHelperTest {
                 .workloadApiClient(workloadApiClient)
                 .build();
 
-        // run KeyStoreHelper
+        try (val keystoreHelper = KeyStoreHelper.create(options)) {
-        new KeyStoreHelper(options);
+            keystoreHelper.run(false);
+        } catch (KeyStoreHelperException e) {
+            fail();
+        }
 
         checkPrivateKeyEntry(keyStoreFilePath, keyStorePass, keyPass, KeyStoreType.getDefaultType(), KeyStoreHelper.DEFAULT_ALIAS);
         val authority1Alias = "example.org.0";
@@ -104,7 +123,7 @@ class KeyStoreHelperTest {
     }
 
     @Test
-    void testNewHelper_keyStore_trustStore_same_file_throwsException() throws SocketEndpointAddressException {
+    void testCreateHelper_keyStore_trustStore_same_file_throwsException() throws SocketEndpointAddressException {
 
         val keyStorefileName = RandomStringUtils.randomAlphabetic(10);
         keyStoreFilePath = Paths.get(keyStorefileName);
@@ -126,22 +145,121 @@ class KeyStoreHelperTest {
                 .build();
 
         try {
-            new KeyStoreHelper(options);
+            KeyStoreHelper.create(options);
             fail("expected exception: KeyStore and TrustStore should use different files");
-        } catch (KeyStoreException e) {
+        } catch (KeyStoreHelperException e) {
             assertEquals("KeyStore and TrustStore should use different files", e.getMessage());
         }
-
     }
 
-    @SneakyThrows
+    @Test
-    @AfterEach
+    void testCreateHelper_keyStore_cannotStoreCerts() throws SocketEndpointAddressException {
-    void tearDown() {
+
-        deleteFile(keyStoreFilePath);
+        keyStoreFilePath = Paths.get("dummy://invalid");
-        deleteFile(trustStoreFilePath);
+        trustStoreFilePath = Paths.get("dummy://otherinvalid");
-        workloadApiClient.close();
+
+        val trustStorePass = "truststore123";
+        val keyStorePass = "keystore123";
+        val keyPass = "keypass123";
+        val alias = "other_alias";
+
+        final KeyStoreHelper.KeyStoreOptions options = KeyStoreHelper.KeyStoreOptions
+                .builder()
+                .keyStorePath(keyStoreFilePath)
+                .keyStorePass(keyStorePass)
+                .trustStorePath(trustStoreFilePath)
+                .trustStorePass(trustStorePass)
+                .keyPass(keyPass)
+                .keyAlias(alias)
+                .workloadApiClient(workloadApiClient)
+                .build();
+
+        try {
+            KeyStoreHelper helper = KeyStoreHelper.create(options);
+            helper.run(false);
+            fail();
+        } catch (KeyStoreHelperException e) {
+            assertEquals("Error running KeyStoreHelper", e.getMessage());
+            assertEquals("java.nio.file.NoSuchFileException: dummy:/invalid", e.getCause().getCause().getMessage());
+        }
     }
 
+    @Test
+    void testCreateKeyStoreHelper_createNewClient() throws Exception {
+        TestUtils.setEnvironmentVariable(Address.SOCKET_ENV_VARIABLE, "unix:/tmp/test");
+        val options = getKeyStoreValidOptions(null);
+        try {
+            KeyStoreHelper.create(options);
+        } catch (KeyStoreHelperException e) {
+            fail();
+        }
+    }
+
+    @Test
+    void testCreateKeyStoreHelper_nullParameter() {
+        try {
+            KeyStoreHelper.create(null);
+        } catch (NullPointerException e) {
+            assertEquals("options is marked non-null but is null", e.getMessage());
+        } catch (SocketEndpointAddressException | KeyStoreHelperException e) {
+            fail();
+        }
+    }
+
+    @Test
+    void testCreateKeyStoreHelper_cannotRunIfClosed() throws Exception {
+        val options = getKeyStoreValidOptions(workloadApiClient);
+        try {
+            KeyStoreHelper helper = KeyStoreHelper.create(options);
+            helper.close();
+            helper.run(true);
+        } catch (IllegalStateException e) {
+            assertEquals("KeyStoreHelper is closed", e.getMessage());
+        } catch (KeyStoreHelperException e) {
+            fail();
+        }
+    }
+
+    @Test
+    void testCreateKeyStoreHelper_throwsExceptionWhenNoUpdateCanBeFetched() throws Exception {
+        val options = getKeyStoreValidOptions(workloadApiClientErrorStub);
+        try {
+            KeyStoreHelper helper = KeyStoreHelper.create(options);
+            helper.run(false);
+            fail();
+        } catch (KeyStoreHelperException e) {
+            assertEquals("Error running KeyStoreHelper", e.getMessage());
+        }
+    }
+
+    private KeyStoreHelper.KeyStoreOptions getKeyStoreValidOptions(WorkloadApiClient workloadApiClient) {
+        val keyStorefileName = RandomStringUtils.randomAlphabetic(10);
+        keyStoreFilePath = Paths.get(keyStorefileName);
+        val trustStorefileName = RandomStringUtils.randomAlphabetic(10);
+        trustStoreFilePath = Paths.get(trustStorefileName);
+
+        val trustStorePass = "truststore123";
+        val keyStorePass = "keystore123";
+        val keyPass = "keypass123";
+        val alias = "other_alias";
+
+        val options = KeyStoreHelper.KeyStoreOptions
+                .builder()
+                .keyStorePath(keyStoreFilePath)
+                .keyStorePass(keyStorePass)
+                .trustStorePath(trustStoreFilePath)
+                .trustStorePass(trustStorePass)
+                .keyPass(keyPass)
+                .keyAlias(alias);
+
+        if (workloadApiClient != null) {
+            options.workloadApiClient(workloadApiClient);
+        }
+
+        return options.build();
+    }
+
+
     private void checkPrivateKeyEntry(Path keyStoreFilePath,
                                       String keyStorePassword,
                                       String privateKeyPassword,

java-spiffe-helper/src/test/java/io/spiffe/helper/keystore/KeyStoreTest.java

@@ -13,7 +13,6 @@ import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 import java.io.IOException;
-import java.net.URI;
 import java.net.URISyntaxException;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -25,6 +24,7 @@ import java.security.UnrecoverableKeyException;
 import java.security.cert.CertificateException;
 import java.security.cert.X509Certificate;
 
+import static io.spiffe.helper.utils.TestUtils.toUri;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.fail;
@@ -222,10 +222,6 @@ public class KeyStoreTest {
         assertEquals(SpiffeId.parse("spiffe://example.org"), spiffeId);
     }
 
-    private URI toUri(String path) throws URISyntaxException {
-        return Thread.currentThread().getContextClassLoader().getResource(path).toURI();
-    }
-
     private void deleteFile(Path filePath) {
         try {
             Files.delete(filePath);

java-spiffe-helper/src/test/java/io/spiffe/helper/keystore/WorkloadApiClientErrorStub.java

@@ -0,0 +1,56 @@
+package io.spiffe.helper.keystore;
+
+import io.spiffe.bundle.jwtbundle.JwtBundleSet;
+import io.spiffe.exception.JwtBundleException;
+import io.spiffe.exception.JwtSvidException;
+import io.spiffe.exception.X509ContextException;
+import io.spiffe.spiffeid.SpiffeId;
+import io.spiffe.svid.jwtsvid.JwtSvid;
+import io.spiffe.workloadapi.Watcher;
+import io.spiffe.workloadapi.WorkloadApiClient;
+import io.spiffe.workloadapi.X509Context;
+import lombok.NonNull;
+
+import java.io.IOException;
+
+public class WorkloadApiClientErrorStub implements WorkloadApiClient {
+
+    @Override
+    public X509Context fetchX509Context() throws X509ContextException {
+        throw new X509ContextException("Testing exception");
+    }
+
+    @Override
+    public void watchX509Context(@NonNull final Watcher<X509Context> watcher) {
+        watcher.onError(new X509ContextException("Testing exception"));
+    }
+
+    @Override
+    public JwtSvid fetchJwtSvid(@NonNull final String audience, final String... extraAudience) throws JwtSvidException {
+        throw new JwtSvidException("Testing exception");
+    }
+
+    @Override
+    public JwtSvid fetchJwtSvid(@NonNull final SpiffeId subject, @NonNull final String audience, final String... extraAudience) throws JwtSvidException {
+        throw new JwtSvidException("Testing exception");
+    }
+
+    @Override
+    public JwtBundleSet fetchJwtBundles() throws JwtBundleException {
+        throw new JwtBundleException("Testing exception");
+    }
+
+    @Override
+    public JwtSvid validateJwtSvid(@NonNull final String token, @NonNull final String audience) throws JwtSvidException {
+        return null;
+    }
+
+    @Override
+    public void watchJwtBundles(@NonNull final Watcher<JwtBundleSet> watcher) {
+        watcher.onError(new JwtBundleException("Testing exception"));
+    }
+
+    @Override
+    public void close() throws IOException {
+    }
+}

java-spiffe-helper/src/test/java/io/spiffe/helper/utils/TestUtils.java

@@ -0,0 +1,53 @@
+package io.spiffe.helper.utils;
+
+import java.lang.reflect.Field;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.Map;
+
+/**
+ * Util methods for testing.
+ */
+public class TestUtils {
+
+    private TestUtils() {
+        throw new UnsupportedOperationException("This is a utility class and cannot be instantiated");
+    }
+
+    public static void setEnvironmentVariable(String variableName, String value) throws Exception {
+        Class<?> processEnvironment = Class.forName("java.lang.ProcessEnvironment");
+
+        Field unmodifiableMapField = getField(processEnvironment, "theUnmodifiableEnvironment");
+        Object unmodifiableMap = unmodifiableMapField.get(null);
+        injectIntoUnmodifiableMap(variableName, value, unmodifiableMap);
+
+        Field mapField = getField(processEnvironment, "theEnvironment");
+        Map<String, String> map = (Map<String, String>) mapField.get(null);
+        map.put(variableName, value);
+    }
+
+    public static Object invokeMethod(Class<?> clazz, String methodName, Object... args) throws NoSuchMethodException, InvocationTargetException, IllegalAccessException {
+        Method method = clazz.getDeclaredMethod(methodName);
+        method.setAccessible(true);
+        return method.invoke(args);
+    }
+
+    public static Field getField(Class<?> clazz, String fieldName) throws NoSuchFieldException {
+        Field field = clazz.getDeclaredField(fieldName);
+        field.setAccessible(true);
+        return field;
+    }
+
+    private static void injectIntoUnmodifiableMap(String key, String value, Object map) throws ReflectiveOperationException {
+        Class unmodifiableMap = Class.forName("java.util.Collections$UnmodifiableMap");
+        Field field = getField(unmodifiableMap, "m");
+        Object obj = field.get(map);
+        ((Map<String, String>) obj).put(key, value);
+    }
+
+    public static URI toUri(String path) throws URISyntaxException {
+        return Thread.currentThread().getContextClassLoader().getResource(path).toURI();
+    }
+}