Add simple test of ZkClientConfigBuilder

2 files changed, 85 insertions, 2 deletions

diff --git a/zookeeper-client-common/src/main/java/com/yahoo/vespa/zookeeper/client/ZkClientConfigBuilder.java b/zookeeper-client-common/src/main/java/com/yahoo/vespa/zookeeper/client/ZkClientConfigBuilder.java
index af75bd7aa5c..0320a76294a 100644
--- a/zookeeper-client-common/src/main/java/com/yahoo/vespa/zookeeper/client/ZkClientConfigBuilder.java
+++ b/zookeeper-client-common/src/main/java/com/yahoo/vespa/zookeeper/client/ZkClientConfigBuilder.java
@@ -30,9 +30,17 @@ public class ZkClientConfigBuilder {
     public static final String SSL_ENABLED_CIPHERSUITES_PROPERTY = "zookeeper.ssl.ciphersuites";
     public static final String SSL_CLIENTAUTH_PROPERTY = "zookeeper.ssl.clientAuth";
 
-    private static final TlsContext tlsContext = getTlsContext().orElse(null);
+    private static final TlsContext defaultTlsContext = getTlsContext().orElse(null);
 
-    public ZkClientConfigBuilder() {}
+    private final TlsContext tlsContext;
+
+    public ZkClientConfigBuilder() {
+         this(defaultTlsContext);
+    }
+
+    public ZkClientConfigBuilder(TlsContext tlsContext) {
+        this.tlsContext = tlsContext;
+    }
 
     public ZKClientConfig toConfig(Path configFile) throws IOException, QuorumPeerConfig.ConfigException {
         String configString = toConfigString();
diff --git a/zookeeper-client-common/src/test/java/com/yahoo/vespa/zookeeper/client/ZkClientConfigBuilderTest.java b/zookeeper-client-common/src/test/java/com/yahoo/vespa/zookeeper/client/ZkClientConfigBuilderTest.java
new file mode 100644
index 00000000000..6b4fb593cef
--- /dev/null
+++ b/zookeeper-client-common/src/test/java/com/yahoo/vespa/zookeeper/client/ZkClientConfigBuilderTest.java
@@ -0,0 +1,75 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.zookeeper.client;
+
+import com.yahoo.security.tls.TlsContext;
+import org.apache.zookeeper.client.ZKClientConfig;
+import org.junit.Test;
+
+import javax.net.ssl.SSLContext;
+import javax.net.ssl.SSLEngine;
+import javax.net.ssl.SSLParameters;
+
+import java.util.List;
+
+import static com.yahoo.vespa.zookeeper.client.ZkClientConfigBuilder.CLIENT_SECURE_PROPERTY;
+import static com.yahoo.vespa.zookeeper.client.ZkClientConfigBuilder.SSL_CLIENTAUTH_PROPERTY;
+import static com.yahoo.vespa.zookeeper.client.ZkClientConfigBuilder.SSL_CONTEXT_SUPPLIER_CLASS_PROPERTY;
+import static com.yahoo.vespa.zookeeper.client.ZkClientConfigBuilder.SSL_ENABLED_CIPHERSUITES_PROPERTY;
+import static com.yahoo.vespa.zookeeper.client.ZkClientConfigBuilder.SSL_ENABLED_PROTOCOLS_PROPERTY;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+
+/**
+ * Tests the zookeeper client config builder.
+ */
+public class ZkClientConfigBuilderTest {
+
+    @Test
+    public void config_when_not_using_tls_context() {
+        ZkClientConfigBuilder builder = new ZkClientConfigBuilder(null);
+        ZKClientConfig config = builder.toConfig();
+        assertEquals("false", config.getProperty(CLIENT_SECURE_PROPERTY));
+        assertNull(config.getProperty(SSL_CONTEXT_SUPPLIER_CLASS_PROPERTY));
+        assertNull(config.getProperty(SSL_CLIENTAUTH_PROPERTY));
+    }
+
+    @Test
+    public void config_when_using_system_tls_context() {
+        ZkClientConfigBuilder builder = new ZkClientConfigBuilder(new MockTlsContext());
+        ZKClientConfig config = builder.toConfig();
+        assertEquals("true", config.getProperty(CLIENT_SECURE_PROPERTY));
+        assertEquals(com.yahoo.vespa.zookeeper.client.VespaSslContextProvider.class.getName(), config.getProperty(SSL_CONTEXT_SUPPLIER_CLASS_PROPERTY));
+        assertEquals("TLSv1.3", config.getProperty(SSL_ENABLED_PROTOCOLS_PROPERTY));
+        assertEquals("TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384", config.getProperty(SSL_ENABLED_CIPHERSUITES_PROPERTY));
+        assertEquals("NEED", config.getProperty(SSL_CLIENTAUTH_PROPERTY));
+    }
+
+    private static class MockTlsContext implements TlsContext {
+
+        @Override
+        public SSLContext context() {
+            return null;
+        }
+
+        @Override
+        public SSLParameters parameters() {
+            SSLParameters parameters = new SSLParameters();
+            parameters.setProtocols(List.of("TLSv1.3").toArray(new String[0]));
+            parameters.setCipherSuites(List.of("TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384").toArray(new String[0]));
+            parameters.setNeedClientAuth(true);
+            return parameters;
+        }
+
+        @Override
+        public SSLEngine createSslEngine() {
+            return null;
+        }
+
+        @Override
+        public SSLEngine createSslEngine(String peerHost, int peerPort) {
+            return null;
+        }
+    }
+
+
+}