Merge pull request searchbox-io#567 from logzio/support-cat-allocations

Support cat allocations

jest-common/src/main/java/io/searchbox/action/AbstractMultiINodeActionBuilder.java

@@ -11,7 +11,7 @@
  */
 @SuppressWarnings("unchecked")
 public abstract class AbstractMultiINodeActionBuilder<T extends Action, K> extends AbstractAction.Builder<T, K> {
-    private Collection<String> nodes = new LinkedList<String>();
+    protected Collection<String> nodes = new LinkedList<String>();
 
     /**
      * Most cluster level APIs allow to specify which nodes to execute on (for example, getting the node stats for a node).

jest-common/src/main/java/io/searchbox/core/Cat.java

@@ -8,8 +8,13 @@
 import com.google.gson.JsonSyntaxException;
 
 import io.searchbox.action.AbstractAction;
+import io.searchbox.action.AbstractMultiINodeActionBuilder;
 import io.searchbox.action.AbstractMultiIndexActionBuilder;
 import io.searchbox.action.AbstractMultiTypeActionBuilder;
+import io.searchbox.strings.StringUtils;
+
+import java.io.UnsupportedEncodingException;
+import java.net.URLEncoder;
 
 /**
  * @author Bartosz Polnik
@@ -27,6 +32,14 @@ public class Cat extends AbstractAction<CatResult> {
     @Override
     protected String buildURI() {
         String uriSuffix = super.buildURI();
+        try {
+            if (!StringUtils.isBlank(nodes)) {
+                uriSuffix += URLEncoder.encode(nodes, CHARSET);
+            }
+        } catch (UnsupportedEncodingException e) {
+            log.error("Error occurred while adding nodes to uri", e);
+        }
+
         return "_cat/" + this.operationPath + (uriSuffix.isEmpty() ? "" : "/") + uriSuffix;
     }
 
@@ -166,6 +179,30 @@ public String operationPath() {
         }
     }
 
+    public static class AllocationBuilder extends AbstractMultiINodeActionBuilder<Cat, AllocationBuilder> implements CatBuilder {
+        private static final String operationPath = "allocation";
+
+        public AllocationBuilder() {
+            setHeader("accept", "application/json");
+            setHeader("content-type", "application/json");
+        }
+
+        @Override
+        public Cat build() {
+            return new Cat(this);
+        }
+
+        @Override
+        public String operationPath() {
+            return operationPath;
+        }
+
+        @Override
+        public String getJoinedNodes() {
+            return nodes.isEmpty() ? null : Joiner.on(',').join(nodes); 
+        }
+    }
+
     protected interface CatBuilder {
         String operationPath();
     }

jest-common/src/test/java/io/searchbox/core/CatAllocationBuilderTest.java

@@ -0,0 +1,44 @@
+package io.searchbox.core;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+
+public class CatAllocationBuilderTest {
+    @Test
+    public void shouldSetApplicationJsonHeader() {
+        Cat cat = new Cat.AllocationBuilder().build();
+        assertEquals("application/json", cat.getHeader("accept"));
+        assertEquals("application/json", cat.getHeader("content-type"));
+    }
+
+    @Test
+    public void shouldGenerateValidUriWhenIndexNotGiven() {
+        Cat cat = new Cat.AllocationBuilder().build();
+        assertEquals("_cat/allocation", cat.getURI());
+    }
+
+    @Test
+    public void shouldGenerateValidUriWhenSingleNodeGiven() {
+        Cat cat = new Cat.AllocationBuilder().addNode("testNode").build();
+        assertEquals("_cat/allocation/testNode", cat.getURI());
+    }
+
+    @Test
+    public void shouldGenerateValidUriWhenNodesGiven() {
+        Cat cat = new Cat.AllocationBuilder().addNode("testNode1").addNode("testNode2").build();
+        assertEquals("_cat/allocation/testNode1%2CtestNode2", cat.getURI());
+    }
+
+    @Test
+    public void shouldGenerateValidUriWhenParameterGiven() {
+        Cat cat = new Cat.AllocationBuilder().setParameter("v", "true").build();
+        assertEquals("_cat/allocation?v=true", cat.getURI());
+    }
+
+    @Test
+    public void shouldGenerateValidUriWhenHeadersParameterGiven() {
+        Cat cat = new Cat.AllocationBuilder().setParameter("h", "shards,disk.indices,disk.used").build();
+        assertEquals("_cat/allocation?h=shards%2Cdisk.indices%2Cdisk.used", cat.getURI());
+    }
+}