diff --git a/server/src/test/java/io/druid/client/selector/TierSelectorStrategyTest.java b/server/src/test/java/io/druid/client/selector/TierSelectorStrategyTest.java new file mode 100644 index 00000000000..956493c41d3 --- /dev/null +++ b/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 getPriorities() + { + return Arrays.asList(2, 0, -1, 1); + } + } + ), + Arrays.asList(lowPriority, mediumPriority, highPriority), + mediumPriority + ); + } + + private void testTierSelectorStrategy( + TierSelectorStrategy tierSelectorStrategy, + List 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.newHashMap(), + Lists.newArrayList(), + Lists.newArrayList(), + new NoneShardSpec(), + 0, + 0L + ), + tierSelectorStrategy + ); + for (QueryableDruidServer server : servers) { + serverSelector.addServer(server); + } + Assert.assertEquals(expectedSelection, serverSelector.pick()); + } + +}