add ut for TierSelectorStrategies

fix test

server/src/test/java/io/druid/client/selector/TierSelectorStrategyTest.java

@@ -0,0 +1,136 @@
+/*
+ * Druid - a distributed column store.
+ * Copyright 2012 - 2015 Metamarkets Group Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.druid.client.selector;
+
+import com.google.common.collect.Lists;
+import io.druid.client.DirectDruidClient;
+import io.druid.client.DruidServer;
+import io.druid.timeline.DataSegment;
+import io.druid.timeline.partition.NoneShardSpec;
+import junit.framework.Assert;
+import org.easymock.EasyMock;
+import org.joda.time.DateTime;
+import org.joda.time.Interval;
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.List;
+
+public class TierSelectorStrategyTest
+{
+
+  @Test
+  public void testHighestPriorityTierSelectorStrategy()
+  {
+    DirectDruidClient client = EasyMock.createMock(DirectDruidClient.class);
+    QueryableDruidServer lowPriority = new QueryableDruidServer(
+        new DruidServer("test1", "localhost", 0, "historical", DruidServer.DEFAULT_TIER, 0),
+        client
+    );
+    QueryableDruidServer highPriority = new QueryableDruidServer(
+        new DruidServer("test1", "localhost", 0, "historical", DruidServer.DEFAULT_TIER, 1),
+        client
+    );
+
+    testTierSelectorStrategy(
+        new HighestPriorityTierSelectorStrategy(new ConnectionCountServerSelectorStrategy()),
+        Arrays.asList(lowPriority, highPriority),
+        highPriority
+    );
+  }
+
+  @Test
+  public void testLowestPriorityTierSelectorStrategy()
+  {
+    DirectDruidClient client = EasyMock.createMock(DirectDruidClient.class);
+    QueryableDruidServer lowPriority = new QueryableDruidServer(
+        new DruidServer("test1", "localhost", 0, "historical", DruidServer.DEFAULT_TIER, 0),
+        client
+    );
+    QueryableDruidServer highPriority = new QueryableDruidServer(
+        new DruidServer("test1", "localhost", 0, "historical", DruidServer.DEFAULT_TIER, 1),
+        client
+    );
+
+    testTierSelectorStrategy(
+        new LowestPriorityTierSelectorStrategy(new ConnectionCountServerSelectorStrategy()),
+        Arrays.asList(lowPriority, highPriority),
+        lowPriority
+    );
+  }
+
+  @Test
+  public void testCustomPriorityTierSelectorStrategy()
+  {
+    DirectDruidClient client = EasyMock.createMock(DirectDruidClient.class);
+    QueryableDruidServer lowPriority = new QueryableDruidServer(
+        new DruidServer("test1", "localhost", 0, "historical", DruidServer.DEFAULT_TIER, -1),
+        client
+    );
+    QueryableDruidServer mediumPriority = new QueryableDruidServer(
+        new DruidServer("test1", "localhost", 0, "historical", DruidServer.DEFAULT_TIER, 0),
+        client
+    );
+    QueryableDruidServer highPriority = new QueryableDruidServer(
+        new DruidServer("test1", "localhost", 0, "historical", DruidServer.DEFAULT_TIER, 1),
+        client
+    );
+
+    testTierSelectorStrategy(
+        new CustomTierSelectorStrategy(
+            new ConnectionCountServerSelectorStrategy(), new CustomTierSelectorStrategyConfig()
+        {
+          @Override
+          public List<Integer> getPriorities()
+          {
+            return Arrays.asList(2, 0, -1, 1);
+          }
+        }
+        ),
+        Arrays.asList(lowPriority, mediumPriority, highPriority),
+        mediumPriority
+    );
+  }
+
+  private void testTierSelectorStrategy(
+      TierSelectorStrategy tierSelectorStrategy,
+      List<QueryableDruidServer> servers,
+      QueryableDruidServer expectedSelection
+  )
+  {
+    final ServerSelector serverSelector = new ServerSelector(
+        new DataSegment(
+            "test",
+            new Interval("2013-01-01/2013-01-02"),
+            new DateTime("2013-01-01").toString(),
+            com.google.common.collect.Maps.<String, Object>newHashMap(),
+            Lists.<String>newArrayList(),
+            Lists.<String>newArrayList(),
+            new NoneShardSpec(),
+            0,
+            0L
+        ),
+        tierSelectorStrategy
+    );
+    for (QueryableDruidServer server : servers) {
+      serverSelector.addServer(server);
+    }
+    Assert.assertEquals(expectedSelection, serverSelector.pick());
+  }
+
+}