NodeChecker factory method added to JestClientFactory - fixes searchb…

…ox-io#312
ozlevka · Mar 27, 2016 · 45739ed · 45739ed
1 parent 4b16bc9
commit 45739ed
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 1 deletion.
diff --git a/jest/src/main/java/io/searchbox/client/JestClientFactory.java b/jest/src/main/java/io/searchbox/client/JestClientFactory.java
@@ -67,7 +67,7 @@ public JestClient getObject() {
         // set discovery (should be set after setting the httpClient on jestClient)
         if (httpClientConfig.isDiscoveryEnabled()) {
             log.info("Node Discovery enabled...");
-            NodeChecker nodeChecker = new NodeChecker(client, httpClientConfig);
+            NodeChecker nodeChecker = createNodeChecker(client, httpClientConfig);
             client.setNodeChecker(nodeChecker);
             nodeChecker.startAsync();
             nodeChecker.awaitRunning();
@@ -226,4 +226,10 @@ protected HttpClientConnectionManager getConnectionManager() {
 
         return retval;
     }
+
+    // Extension point
+    protected NodeChecker createNodeChecker(JestHttpClient client, HttpClientConfig httpClientConfig) {
+        return new NodeChecker(client, httpClientConfig);
+    }
+
 }
diff --git a/jest/src/test/java/io/searchbox/client/JestClientFactoryTest.java b/jest/src/test/java/io/searchbox/client/JestClientFactoryTest.java
@@ -1,6 +1,8 @@
 package io.searchbox.client;
 
+import io.searchbox.client.config.ClientConfig;
 import io.searchbox.client.config.HttpClientConfig;
+import io.searchbox.client.config.discovery.NodeChecker;
 import io.searchbox.client.http.JestHttpClient;
 import org.apache.http.HttpHost;
 import org.apache.http.client.config.RequestConfig;
@@ -11,6 +13,7 @@
 import org.apache.http.impl.nio.conn.PoolingNHttpClientConnectionManager;
 import org.apache.http.nio.conn.NHttpClientConnectionManager;
 import org.junit.Test;
+import org.mockito.Mockito;
 
 import static org.junit.Assert.*;
 
@@ -89,4 +92,34 @@ public void multiThreadedClientCreation() {
         assertEquals(5, ((PoolingNHttpClientConnectionManager) nConnectionManager).getMaxPerRoute(routeOne));
         assertEquals(6, ((PoolingNHttpClientConnectionManager) nConnectionManager).getMaxPerRoute(routeTwo));
     }
+
+    @Test
+    public void clientCreationWithDiscoveryAndOverridenNodeChecker() {
+        JestClientFactory factory = Mockito.spy(new ExtendedJestClientFactory());
+        HttpClientConfig httpClientConfig = Mockito.spy(new HttpClientConfig.Builder("http://localhost:9200")
+                .discoveryEnabled(true)
+                .build());
+        factory.setHttpClientConfig(httpClientConfig);
+        JestHttpClient jestClient = (JestHttpClient) factory.getObject();
+        assertTrue(jestClient != null);
+        assertNotNull(jestClient.getAsyncClient());
+        assertEquals(jestClient.getServerPoolSize(), 1);
+        assertEquals("server list should contain localhost:9200",
+                "http://localhost:9200", jestClient.getNextServer());
+        Mockito.verify(factory, Mockito.times(1)).createNodeChecker(Mockito.any(JestHttpClient.class),
+                                                                    Mockito.same(httpClientConfig));
+    }
+
+    class ExtendedJestClientFactory extends JestClientFactory {
+        @Override
+        protected NodeChecker createNodeChecker(JestHttpClient client, HttpClientConfig httpClientConfig) {
+            return new OtherNodeChecker(client, httpClientConfig);
+        }
+    }
+
+    class OtherNodeChecker extends NodeChecker {
+        public OtherNodeChecker(JestClient jestClient, ClientConfig clientConfig) {
+            super(jestClient, clientConfig);
+        }
+    }
 }