diff options
Diffstat (limited to 'exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr')
9 files changed, 1381 insertions, 0 deletions
diff --git a/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/TestRMConfigLoad.java b/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/TestRMConfigLoad.java new file mode 100644 index 000000000..16f7f64f0 --- /dev/null +++ b/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/TestRMConfigLoad.java @@ -0,0 +1,117 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.drill.exec.resourcemgr; + +import org.apache.drill.categories.ResourceManagerTest; +import org.apache.drill.exec.ExecConstants; +import org.apache.drill.exec.resourcemgr.config.QueryQueueConfig; +import org.apache.drill.exec.resourcemgr.config.RMCommonDefaults; +import org.apache.drill.exec.resourcemgr.config.ResourcePoolTree; +import org.apache.drill.exec.resourcemgr.config.selectors.AclSelector; +import org.apache.drill.exec.work.foreman.rm.DefaultResourceManager; +import org.apache.drill.exec.work.foreman.rm.DistributedResourceManager; +import org.apache.drill.exec.work.foreman.rm.ResourceManager; +import org.apache.drill.test.BaseDirTestWatcher; +import org.apache.drill.test.ClusterFixture; +import org.apache.drill.test.ClusterFixtureBuilder; +import org.apache.drill.test.DrillTest; +import org.junit.Ignore; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +@Ignore("These tests will be ignored until integration with new DistributedResourceManager is done") +@Category(ResourceManagerTest.class) +public final class TestRMConfigLoad extends DrillTest { + + @Rule + public final BaseDirTestWatcher dirTestWatcher = new BaseDirTestWatcher(); + + @Test + public void testDefaultRMConfig() throws Exception { + ClusterFixtureBuilder fixtureBuilder = ClusterFixture.builder(dirTestWatcher) + .configProperty(ExecConstants.RM_ENABLED, true) + .configProperty(ExecConstants.DRILL_PORT_HUNT, true) + .withLocalZk(); + + try (ClusterFixture cluster = fixtureBuilder.build()) { + ResourceManager resourceManager = cluster.drillbit().getContext().getResourceManager(); + assertTrue(resourceManager instanceof DistributedResourceManager); + + ResourcePoolTree poolTree = ((DistributedResourceManager) resourceManager).getRmPoolTree(); + assertTrue("In drill-rm-default root pool is not leaf pool", poolTree.getRootPool().isLeafPool()); + assertTrue("selector in drill-rm-default is not acl selector", + poolTree.getRootPool().getSelector() instanceof AclSelector); + assertEquals("max_query_memory_per_node in drill-rm-default is not configured with expected default value", + 8 * 1024L, poolTree.getRootPool().getMaxQueryMemoryPerNode()); + assertEquals("queue_selection_policy in drill-rm-default is not configured with expected default value", + RMCommonDefaults.ROOT_POOL_DEFAULT_QUEUE_SELECTION_POLICY, + poolTree.getSelectionPolicyInUse().getSelectionPolicy()); + assertEquals("memory share of root pool in drill-rm-default is not configured with expected default value", + RMCommonDefaults.ROOT_POOL_DEFAULT_MEMORY_PERCENT, poolTree.getResourceShare(), 0); + + final QueryQueueConfig defaultQueue = poolTree.getRootPool().getQueryQueue(); + assertEquals("max_admissible in drill-rm-default is not configured with expected default value", + RMCommonDefaults.MAX_ADMISSIBLE_QUERY_COUNT, defaultQueue.getMaxAdmissibleQueries()); + assertEquals("max_waiting in drill-rm-default is not configured with expected default value", + RMCommonDefaults.MAX_WAITING_QUERY_COUNT, defaultQueue.getMaxWaitingQueries()); + assertEquals("max_wait_timeout in drill-rm-default is not configured with expected default value", + RMCommonDefaults.MAX_WAIT_TIMEOUT_IN_MS, defaultQueue.getWaitTimeoutInMs()); + assertEquals("wait_for_preferred_nodes in drill-rm-default is not configured with expected default value", + RMCommonDefaults.WAIT_FOR_PREFERRED_NODES, defaultQueue.waitForPreferredNodes()); + } + } + + @Test + public void testDefaultRMWithLocalCoordinatorAndRMEnabled() throws Exception { + ClusterFixtureBuilder fixtureBuilder = ClusterFixture.builder(dirTestWatcher) + .configProperty(ExecConstants.RM_ENABLED, true); + + try (ClusterFixture cluster = fixtureBuilder.build()) { + ResourceManager resourceManager = cluster.drillbit().getContext().getResourceManager(); + assertTrue(resourceManager instanceof DefaultResourceManager); + } + } + + @Test + public void testDefaultRMWithLocalCoordinatorAndRMDisabled() throws Exception { + ClusterFixtureBuilder fixtureBuilder = ClusterFixture.builder(dirTestWatcher) + .configProperty(ExecConstants.RM_ENABLED, false); + + try (ClusterFixture cluster = fixtureBuilder.build()) { + ResourceManager resourceManager = cluster.drillbit().getContext().getResourceManager(); + assertTrue(resourceManager instanceof DefaultResourceManager); + } + } + + @Test + public void testDefaultRMOnlyRMDisabled() throws Exception { + ClusterFixtureBuilder fixtureBuilder = ClusterFixture.builder(dirTestWatcher) + .configProperty(ExecConstants.RM_ENABLED, false) + .configProperty(ExecConstants.DRILL_PORT_HUNT, true) + .withLocalZk(); + + try (ClusterFixture cluster = fixtureBuilder.build()) { + ResourceManager resourceManager = cluster.drillbit().getContext().getResourceManager(); + assertTrue(resourceManager instanceof DefaultResourceManager); + } + } +} diff --git a/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/TestResourcePoolTree.java b/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/TestResourcePoolTree.java new file mode 100644 index 000000000..fea132d41 --- /dev/null +++ b/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/TestResourcePoolTree.java @@ -0,0 +1,277 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.drill.exec.resourcemgr; + +import com.typesafe.config.Config; +import com.typesafe.config.ConfigFactory; +import com.typesafe.config.ConfigValueFactory; +import org.apache.drill.categories.ResourceManagerTest; +import org.apache.drill.exec.ExecConstants; +import org.apache.drill.exec.ops.QueryContext; +import org.apache.drill.exec.resourcemgr.config.QueueAssignmentResult; +import org.apache.drill.exec.resourcemgr.config.ResourcePool; +import org.apache.drill.exec.resourcemgr.config.ResourcePoolImpl; +import org.apache.drill.exec.resourcemgr.config.ResourcePoolTree; +import org.apache.drill.exec.resourcemgr.config.ResourcePoolTreeImpl; +import org.apache.drill.exec.resourcemgr.config.exception.RMConfigException; +import org.apache.drill.exec.server.options.OptionValue; +import org.junit.After; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static org.apache.drill.exec.resourcemgr.config.RMCommonDefaults.ROOT_POOL_DEFAULT_QUEUE_SELECTION_POLICY; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +@Category(ResourceManagerTest.class) +public final class TestResourcePoolTree { + + private static final Map<String, Object> poolTreeConfig = new HashMap<>(); + + private static final Map<String, Object> pool1 = new HashMap<>(); + + private static final Map<String, Object> pool2 = new HashMap<>(); + + private static final Map<String, Object> queue1 = new HashMap<>(); + + private static final List<Object> childResourcePools = new ArrayList<>(); + + private static final Map<String, Object> tagSelectorConfig1 = new HashMap<>(); + + private static final Map<String, Object> tagSelectorConfig2 = new HashMap<>(); + + private static final QueryContext mockContext = mock(QueryContext.class); + + @BeforeClass + public static void testSuiteSetup() { + pool1.put(ResourcePoolImpl.POOL_NAME_KEY, "dev"); + pool1.put(ResourcePoolImpl.POOL_MEMORY_SHARE_KEY, 0.80); + + pool2.put(ResourcePoolImpl.POOL_NAME_KEY, "qa"); + pool2.put(ResourcePoolImpl.POOL_MEMORY_SHARE_KEY, 0.20); + + queue1.put("max_query_memory_per_node", 5534); + + tagSelectorConfig1.put("tag", "small"); + tagSelectorConfig2.put("tag", "large"); + } + + @After + public void afterTestCleanup() { + // cleanup resource tree + poolTreeConfig.clear(); + + // cleanup pools + pool1.remove(ResourcePoolImpl.POOL_QUEUE_KEY); + pool1.remove(ResourcePoolImpl.POOL_SELECTOR_KEY); + + pool2.remove(ResourcePoolImpl.POOL_QUEUE_KEY); + pool2.remove(ResourcePoolImpl.POOL_SELECTOR_KEY); + + childResourcePools.clear(); + } + + private ResourcePoolTree getPoolTreeConfig() throws RMConfigException { + poolTreeConfig.put(ResourcePoolImpl.POOL_NAME_KEY, "drill"); + poolTreeConfig.put(ResourcePoolImpl.POOL_CHILDREN_POOLS_KEY, childResourcePools); + + Config rmConfig = ConfigFactory.empty() + .withValue("drill.exec.rm", ConfigValueFactory.fromMap(poolTreeConfig)); + return new ResourcePoolTreeImpl(rmConfig, 10000, 10, 2); + } + + private boolean checkExpectedVsActualPools(List<ResourcePool> actual, List<String> expectedNames) { + if (actual.size() != expectedNames.size()) { + return false; + } + + for (ResourcePool pool : actual) { + if (!expectedNames.contains(pool.getPoolName())) { + return false; + } + } + return true; + } + + @Test + public void testTreeWith2LeafPool() throws Exception { + // pool with tag selector + pool1.put(ResourcePoolImpl.POOL_QUEUE_KEY, queue1); + pool1.put(ResourcePoolImpl.POOL_SELECTOR_KEY, tagSelectorConfig1); + + // pool with default selector + pool2.put(ResourcePoolImpl.POOL_QUEUE_KEY, queue1); + + childResourcePools.add(pool1); + childResourcePools.add(pool2); + + ResourcePoolTree configTree = getPoolTreeConfig(); + + // get all leaf queues names + Set<String> expectedLeafQueue = new HashSet<>(); + expectedLeafQueue.add((String)pool1.get("pool_name")); + expectedLeafQueue.add((String)pool2.get("pool_name")); + + assertEquals("Root pool is different than expected", "drill", configTree.getRootPool().getPoolName()); + assertEquals("Expected and actual leaf queue names are different", expectedLeafQueue, + configTree.getAllLeafQueues().keySet()); + assertEquals("Unexpected Selection policy is in use", ROOT_POOL_DEFAULT_QUEUE_SELECTION_POLICY, + configTree.getSelectionPolicyInUse().getSelectionPolicy()); + } + + @Test(expected = RMConfigException.class) + public void testDuplicateLeafPool() throws Exception { + // leaf pool + pool1.put(ResourcePoolImpl.POOL_QUEUE_KEY, queue1); + childResourcePools.add(pool1); + childResourcePools.add(pool1); + + getPoolTreeConfig(); + } + + @Test(expected = RMConfigException.class) + public void testMissingQueueAtLeafPool() throws Exception { + // leaf pool with queue + pool1.put(ResourcePoolImpl.POOL_QUEUE_KEY, queue1); + pool1.put(ResourcePoolImpl.POOL_SELECTOR_KEY, tagSelectorConfig1); + childResourcePools.add(pool1); + childResourcePools.add(pool2); + + getPoolTreeConfig(); + } + + @Test(expected = RMConfigException.class) + public void testInvalidQueueAtLeafPool() throws Exception { + // leaf pool with invalid queue + int initialValue = (Integer)queue1.remove("max_query_memory_per_node"); + + try { + pool1.put(ResourcePoolImpl.POOL_QUEUE_KEY, queue1); + pool1.put(ResourcePoolImpl.POOL_SELECTOR_KEY, tagSelectorConfig1); + childResourcePools.add(pool1); + + getPoolTreeConfig(); + } finally { + queue1.put("max_query_memory_per_node", initialValue); + } + } + + @Test + public void testRootPoolAsLeaf() throws Exception { + // leaf pool with queue + poolTreeConfig.put(ResourcePoolImpl.POOL_NAME_KEY, "drill"); + poolTreeConfig.put(ResourcePoolImpl.POOL_QUEUE_KEY, queue1); + poolTreeConfig.put(ResourcePoolImpl.POOL_SELECTOR_KEY, tagSelectorConfig1); + + Config rmConfig = ConfigFactory.empty() + .withValue("drill.exec.rm", ConfigValueFactory.fromMap(poolTreeConfig)); + ResourcePoolTree poolTree = new ResourcePoolTreeImpl(rmConfig, 10000, 10, 2); + + assertTrue("Root pool is not a leaf pool", poolTree.getRootPool().isLeafPool()); + assertEquals("Root pool name is not drill", "drill", poolTree.getRootPool().getPoolName()); + assertTrue("Root pool is not the only leaf pool", poolTree.getAllLeafQueues().size() == 1); + assertTrue("Root pool name is not same as leaf pool name", poolTree.getAllLeafQueues().containsKey("drill")); + assertFalse("Root pool should not be a default pool", poolTree.getRootPool().isDefaultPool()); + } + + @Test + public void testTreeWithLeafAndIntermediatePool() throws Exception { + // left leaf pool1 with tag selector + pool1.put(ResourcePoolImpl.POOL_QUEUE_KEY, queue1); + pool1.put(ResourcePoolImpl.POOL_SELECTOR_KEY, tagSelectorConfig1); + + // left leaf pool2 with default selector + pool2.put(ResourcePoolImpl.POOL_QUEUE_KEY, queue1); + + // intermediate left pool1 with 2 leaf pools (pool1, pool2) + Map<String, Object> interPool1 = new HashMap<>(); + List<Object> childPools1 = new ArrayList<>(); + childPools1.add(pool1); + childPools1.add(pool2); + interPool1.put(ResourcePoolImpl.POOL_NAME_KEY, "eng"); + interPool1.put(ResourcePoolImpl.POOL_MEMORY_SHARE_KEY, 0.9); + interPool1.put(ResourcePoolImpl.POOL_CHILDREN_POOLS_KEY, childPools1); + + // right leaf pool + Map<String, Object> rightLeafPool = new HashMap<>(); + rightLeafPool.put(ResourcePoolImpl.POOL_NAME_KEY, "marketing"); + rightLeafPool.put(ResourcePoolImpl.POOL_MEMORY_SHARE_KEY, 0.1); + rightLeafPool.put(ResourcePoolImpl.POOL_QUEUE_KEY, queue1); + rightLeafPool.put(ResourcePoolImpl.POOL_SELECTOR_KEY, tagSelectorConfig2); + + childResourcePools.add(interPool1); + childResourcePools.add(rightLeafPool); + ResourcePoolTree configTree = getPoolTreeConfig(); + + // Test successful selection of all leaf pools + OptionValue testOption = OptionValue.create(OptionValue.AccessibleScopes.SESSION_AND_QUERY, ExecConstants + .RM_QUERY_TAGS_KEY, "small,large", OptionValue.OptionScope.SESSION); + when(mockContext.getOption(ExecConstants.RM_QUERY_TAGS_KEY)).thenReturn(testOption); + QueueAssignmentResult assignmentResult = configTree.selectAllQueues(mockContext); + List<ResourcePool> selectedPools = assignmentResult.getSelectedLeafPools(); + List<String> expectedPools = new ArrayList<>(); + expectedPools.add("dev"); + expectedPools.add("qa"); + expectedPools.add("marketing"); + + assertTrue("All leaf pools are not selected", selectedPools.size() == 3); + assertTrue("Selected leaf pools and expected pools are different", + checkExpectedVsActualPools(selectedPools, expectedPools)); + + // Test successful selection of multiple leaf pools + expectedPools.clear(); + testOption = OptionValue.create(OptionValue.AccessibleScopes.SESSION_AND_QUERY, ExecConstants + .RM_QUERY_TAGS_KEY, "small", OptionValue.OptionScope.SESSION); + when(mockContext.getOption(ExecConstants.RM_QUERY_TAGS_KEY)).thenReturn(testOption); + assignmentResult = configTree.selectAllQueues(mockContext); + selectedPools = assignmentResult.getSelectedLeafPools(); + expectedPools.add("qa"); + expectedPools.add("dev"); + assertTrue("Expected 2 pools to be selected", selectedPools.size() == 2); + assertTrue("Selected leaf pools and expected pools are different", + checkExpectedVsActualPools(selectedPools, expectedPools)); + + // Test successful selection of only left default pool + expectedPools.clear(); + testOption = OptionValue.create(OptionValue.AccessibleScopes.SESSION_AND_QUERY, ExecConstants + .RM_QUERY_TAGS_KEY, "medium", OptionValue.OptionScope.SESSION); + when(mockContext.getOption(ExecConstants.RM_QUERY_TAGS_KEY)).thenReturn(testOption); + assignmentResult = configTree.selectAllQueues(mockContext); + selectedPools = assignmentResult.getSelectedLeafPools(); + expectedPools.add("qa"); + assertTrue("More than one leaf pool is selected", selectedPools.size() == 1); + assertTrue("Selected leaf pools and expected pools are different", + checkExpectedVsActualPools(selectedPools, expectedPools)); + + // cleanup + interPool1.clear(); + rightLeafPool.clear(); + expectedPools.clear(); + } +} diff --git a/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/config/selectionpolicy/TestBestFitSelectionPolicy.java b/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/config/selectionpolicy/TestBestFitSelectionPolicy.java new file mode 100644 index 000000000..3dcde00ca --- /dev/null +++ b/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/config/selectionpolicy/TestBestFitSelectionPolicy.java @@ -0,0 +1,122 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.drill.exec.resourcemgr.config.selectionpolicy; + +import org.apache.drill.categories.ResourceManagerTest; +import org.apache.drill.exec.ops.QueryContext; +import org.apache.drill.exec.proto.UserBitShared; +import org.apache.drill.exec.resourcemgr.NodeResources; +import org.apache.drill.exec.resourcemgr.config.QueryQueueConfig; +import org.apache.drill.exec.resourcemgr.config.RMCommonDefaults; +import org.apache.drill.exec.resourcemgr.config.ResourcePool; +import org.apache.drill.exec.resourcemgr.config.exception.QueueSelectionException; +import org.junit.After; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.util.ArrayList; +import java.util.List; + +import static org.junit.Assert.assertEquals; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +@Category(ResourceManagerTest.class) +public final class TestBestFitSelectionPolicy { + + private static QueueSelectionPolicy selectionPolicy; + + private static final QueryContext queryContext = mock(QueryContext.class); + + private static final NodeResources queryMaxResources = new NodeResources(1500*1024L*1024L, 2); + + private static final List<Long> poolMemory = new ArrayList<>(); + + @BeforeClass + public static void testSetup() { + selectionPolicy = QueueSelectionPolicyFactory.createSelectionPolicy + (RMCommonDefaults.ROOT_POOL_DEFAULT_QUEUE_SELECTION_POLICY); + when(queryContext.getQueryId()).thenReturn(UserBitShared.QueryId.getDefaultInstance()); + } + + @After + public void afterTestSetup() { + poolMemory.clear(); + } + + private void testCommonHelper(long expectedPoolMem) throws Exception { + List<ResourcePool> inputPools = new ArrayList<>(); + ResourcePool expectedPool = null; + + for (Long poolMemory : poolMemory) { + final ResourcePool testPool = mock(ResourcePool.class); + final QueryQueueConfig testPoolQueue = mock(QueryQueueConfig.class); + when(testPool.getQueryQueue()).thenReturn(testPoolQueue); + when(testPoolQueue.getMaxQueryMemoryInMBPerNode()).thenReturn(poolMemory); + inputPools.add(testPool); + + if (poolMemory == expectedPoolMem) { + expectedPool = testPool; + } + } + + ResourcePool selectedPool = selectionPolicy.selectQueue(inputPools, queryContext, queryMaxResources); + assertEquals("Selected Pool and expected pool is different", expectedPool, selectedPool); + } + + @Test(expected = QueueSelectionException.class) + public void testWithNoPool() throws Exception { + testCommonHelper(0); + } + + @Test + public void testWithSinglePool() throws Exception { + poolMemory.add(1000L); + testCommonHelper(1000); + } + + @Test + public void testWithMultiplePoolWithGreaterMaxNodeMemory() throws Exception { + poolMemory.add(2500L); + poolMemory.add(2000L); + testCommonHelper(2000); + } + + @Test + public void testWithMultiplePoolWithLesserMaxNodeMemory() throws Exception { + poolMemory.add(700L); + poolMemory.add(500L); + testCommonHelper(700); + } + + @Test + public void testMixOfPoolLess_Greater_MaxNodeMemory() throws Exception { + poolMemory.add(1000L); + poolMemory.add(2000L); + testCommonHelper(2000); + } + + @Test + public void testMixOfPoolLess_Greater_EqualMaxNodeMemory() throws Exception { + poolMemory.add(1000L); + poolMemory.add(2000L); + poolMemory.add(1500L); + testCommonHelper(1500); + } +} diff --git a/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/config/selectionpolicy/TestDefaultSelectionPolicy.java b/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/config/selectionpolicy/TestDefaultSelectionPolicy.java new file mode 100644 index 000000000..a78d106f3 --- /dev/null +++ b/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/config/selectionpolicy/TestDefaultSelectionPolicy.java @@ -0,0 +1,103 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.drill.exec.resourcemgr.config.selectionpolicy; + +import org.apache.drill.categories.ResourceManagerTest; +import org.apache.drill.exec.ops.QueryContext; +import org.apache.drill.exec.proto.UserBitShared; +import org.apache.drill.exec.resourcemgr.config.ResourcePool; +import org.apache.drill.exec.resourcemgr.config.exception.QueueSelectionException; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.util.ArrayList; +import java.util.List; + +import static org.junit.Assert.assertEquals; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +@Category(ResourceManagerTest.class) +public final class TestDefaultSelectionPolicy { + + private static QueueSelectionPolicy selectionPolicy; + + private static final QueryContext queryContext = mock(QueryContext.class); + + @BeforeClass + public static void testSetup() { + selectionPolicy = QueueSelectionPolicyFactory.createSelectionPolicy( + QueueSelectionPolicy.SelectionPolicy.DEFAULT); + when(queryContext.getQueryId()).thenReturn(UserBitShared.QueryId.getDefaultInstance()); + } + + @Test(expected = QueueSelectionException.class) + public void testWithNoDefaultPool() throws Exception { + List<ResourcePool> inputPools = new ArrayList<>(); + final ResourcePool testPool1 = mock(ResourcePool.class); + when(testPool1.isDefaultPool()).thenReturn(false); + final ResourcePool testPool2 = mock(ResourcePool.class); + when(testPool2.isDefaultPool()).thenReturn(false); + + inputPools.add(testPool1); + inputPools.add(testPool2); + selectionPolicy.selectQueue(inputPools, queryContext, null); + } + + @Test + public void testWithSingleDefaultPool() throws Exception { + List<ResourcePool> inputPools = new ArrayList<>(); + final ResourcePool testPool1 = mock(ResourcePool.class); + when(testPool1.isDefaultPool()).thenReturn(true); + inputPools.add(testPool1); + + final ResourcePool selectedPool = selectionPolicy.selectQueue(inputPools, queryContext, null); + assertEquals("Selected Pool and expected pool is different",testPool1, selectedPool); + } + + @Test + public void testWithMultipleDefaultPool() throws Exception { + List<ResourcePool> inputPools = new ArrayList<>(); + final ResourcePool testPool1 = mock(ResourcePool.class); + when(testPool1.isDefaultPool()).thenReturn(true); + final ResourcePool testPool2 = mock(ResourcePool.class); + when(testPool2.isDefaultPool()).thenReturn(true); + + inputPools.add(testPool1); + inputPools.add(testPool2); + + final ResourcePool selectedPool = selectionPolicy.selectQueue(inputPools, queryContext, null); + assertEquals("Selected Pool and expected pool is different",testPool1, selectedPool); + } + + @Test + public void testMixOfDefaultAndNonDefaultPool() throws Exception { + List<ResourcePool> inputPools = new ArrayList<>(); + final ResourcePool testPool1 = mock(ResourcePool.class); + when(testPool1.isDefaultPool()).thenReturn(false); + final ResourcePool testPool2 = mock(ResourcePool.class); + when(testPool2.isDefaultPool()).thenReturn(true); + + inputPools.add(testPool1); + inputPools.add(testPool2); + + final ResourcePool selectedPool = selectionPolicy.selectQueue(inputPools, queryContext, null); + assertEquals("Selected Pool and expected pool is different",testPool2, selectedPool); + } +} diff --git a/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/config/selectors/TestAclSelector.java b/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/config/selectors/TestAclSelector.java new file mode 100644 index 000000000..fbdb108fd --- /dev/null +++ b/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/config/selectors/TestAclSelector.java @@ -0,0 +1,265 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.drill.exec.resourcemgr.config.selectors; + +import com.typesafe.config.Config; +import com.typesafe.config.ConfigFactory; +import com.typesafe.config.ConfigValueFactory; +import org.apache.drill.categories.ResourceManagerTest; +import org.apache.drill.exec.resourcemgr.config.exception.RMConfigException; +import org.junit.After; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +@Category(ResourceManagerTest.class) +public final class TestAclSelector { + + private static final List<String> groupsValue = new ArrayList<>(); + + private static final List<String> usersValue = new ArrayList<>(); + + private static final List<String> emptyList = new ArrayList<>(); + + private static final Map<String, List<String>> aclConfigValue = new HashMap<>(); + + private boolean checkIfSame(Object[] expected, Object[] actual) { + Arrays.sort(expected); + Arrays.sort(actual); + return Arrays.equals(expected, actual); + } + + @After + public void cleanupAfterTest() { + groupsValue.clear(); + usersValue.clear(); + aclConfigValue.clear(); + } + + private ResourcePoolSelector testNegativeHelper() throws RMConfigException { + final Config testConfig = ConfigFactory.empty() + .withValue("acl", ConfigValueFactory.fromMap(aclConfigValue)); + return ResourcePoolSelectorFactory.createSelector(testConfig); + } + + private ResourcePoolSelector testCommonHelper(List<String> expectedPositiveUsers, List<String> expectedPositiveGroups, + List<String> expectedNegativeUsers, List<String> expectedNegativeGroups) + throws RMConfigException { + ResourcePoolSelector testSelector = testNegativeHelper(); + assertTrue("TestSelector is not a ACL selector", testSelector instanceof AclSelector); + assertTrue("Expected +ve users config mismatched with actual config", + checkIfSame(expectedPositiveUsers.toArray(), ((AclSelector) testSelector).getAllowedUsers().toArray())); + assertTrue("Expected +ve groups config mismatched with actual config", + checkIfSame(expectedPositiveGroups.toArray(), ((AclSelector) testSelector).getAllowedGroups().toArray())); + assertTrue("Expected -ve users config mismatched with actual config", + checkIfSame(expectedNegativeUsers.toArray(), ((AclSelector) testSelector).getDeniedUsers().toArray())); + assertTrue("Expected -ve groups config mismatched with actual config", + checkIfSame(expectedNegativeGroups.toArray(), ((AclSelector) testSelector).getDeniedGroups().toArray())); + return testSelector; + } + + @Test + public void testValidACLSelector_shortSyntax() throws Exception { + groupsValue.add("sales"); + groupsValue.add("marketing"); + + usersValue.add("user1"); + usersValue.add("user2"); + + aclConfigValue.put("groups", groupsValue); + aclConfigValue.put("users", usersValue); + AclSelector testSelector = (AclSelector) testCommonHelper(usersValue, groupsValue, emptyList, emptyList); + + // check based on valid/invalid user + Set<String> groups = new HashSet<>(); + assertFalse(testSelector.checkQueryUserGroups("user3", groups)); + assertTrue(testSelector.checkQueryUserGroups("user1", groups)); + + // check based on correct group + groups.add("sales"); + assertTrue(testSelector.checkQueryUserGroups("user3", groups)); + } + + @Test + public void testACLSelector_onlyUsers() throws Exception { + usersValue.add("user1"); + aclConfigValue.put("users", usersValue); + testCommonHelper(usersValue, groupsValue, emptyList, emptyList); + } + + @Test + public void testACLSelector_onlyGroups() throws Exception { + groupsValue.add("group1"); + aclConfigValue.put("groups", groupsValue); + testCommonHelper(usersValue, groupsValue, emptyList, emptyList); + } + + @Test(expected = RMConfigException.class) + public void testInValidACLSelector_shortSyntax() throws Exception { + aclConfigValue.put("groups", new ArrayList<>()); + aclConfigValue.put("users", new ArrayList<>()); + testNegativeHelper(); + } + + @Test + public void testValidACLSelector_longSyntax() throws Exception { + + groupsValue.add("sales:+"); + groupsValue.add("marketing:-"); + + List<String> expectedAllowedGroups = new ArrayList<>(); + expectedAllowedGroups.add("sales"); + + List<String> expectedDisAllowedGroups = new ArrayList<>(); + expectedDisAllowedGroups.add("marketing"); + + usersValue.add("user1:+"); + usersValue.add("user2:-"); + + List<String> expectedAllowedUsers = new ArrayList<>(); + expectedAllowedUsers.add("user1"); + List<String> expectedDisAllowedUsers = new ArrayList<>(); + expectedDisAllowedUsers.add("user2"); + + aclConfigValue.put("groups", groupsValue); + aclConfigValue.put("users", usersValue); + AclSelector testSelector = (AclSelector)testCommonHelper(expectedAllowedUsers, expectedAllowedGroups, + expectedDisAllowedUsers, expectedDisAllowedGroups); + + Set<String> queryGroups = new HashSet<>(); + // Negative user/group + queryGroups.add("marketing"); + assertFalse(testSelector.checkQueryUserGroups("user2", queryGroups)); + + // Invalid user -ve group + assertFalse(testSelector.checkQueryUserGroups("user3", queryGroups)); + + // -ve user +ve group + queryGroups.clear(); + queryGroups.add("sales"); + assertFalse(testSelector.checkQueryUserGroups("user2", queryGroups)); + + // Invalid user +ve group + assertTrue(testSelector.checkQueryUserGroups("user3", queryGroups)); + } + + @Test(expected = RMConfigException.class) + public void testInvalidLongSyntaxIdentifier() throws Exception { + groupsValue.add("sales:|"); + aclConfigValue.put("groups", groupsValue); + testNegativeHelper(); + } + + @Test + public void testMixLongShortAclSyntax() throws Exception { + groupsValue.add("groups1"); + groupsValue.add("groups2:+"); + groupsValue.add("groups3:-"); + + List<String> expectedAllowedGroups = new ArrayList<>(); + expectedAllowedGroups.add("groups1"); + expectedAllowedGroups.add("groups2"); + + List<String> expectedDisAllowedGroups = new ArrayList<>(); + expectedDisAllowedGroups.add("groups3"); + + usersValue.add("user1"); + usersValue.add("user2:+"); + usersValue.add("user3:-"); + + List<String> expectedAllowedUsers = new ArrayList<>(); + expectedAllowedUsers.add("user1"); + expectedAllowedUsers.add("user2"); + + List<String> expectedDisAllowedUsers = new ArrayList<>(); + expectedDisAllowedUsers.add("user3"); + + aclConfigValue.put("groups", groupsValue); + aclConfigValue.put("users", usersValue); + + testCommonHelper(expectedAllowedUsers, expectedAllowedGroups, expectedDisAllowedUsers, expectedDisAllowedGroups); + } + + @Test + public void testSameUserBothInPositiveNegative() throws Exception { + usersValue.add("user1:+"); + usersValue.add("user1:-"); + + List<String> expectedDisAllowedUsers = new ArrayList<>(); + expectedDisAllowedUsers.add("user1"); + + aclConfigValue.put("users", usersValue); + testCommonHelper(emptyList, emptyList, expectedDisAllowedUsers, emptyList); + } + + @Test + public void testStarInPositiveUsers() throws Exception { + usersValue.add("*:+"); + usersValue.add("user1:-"); + + List<String> expectedAllowedUsers = new ArrayList<>(); + expectedAllowedUsers.add("*"); + + List<String> expectedDisAllowedUsers = new ArrayList<>(); + expectedDisAllowedUsers.add("user1"); + + aclConfigValue.put("users", usersValue); + AclSelector testSelector = (AclSelector)testCommonHelper(expectedAllowedUsers, emptyList, + expectedDisAllowedUsers, emptyList); + + Set<String> queryGroups = new HashSet<>(); + // -ve user with Invalid groups + assertFalse(testSelector.checkQueryUserGroups("user1", queryGroups)); + + // Other user with invalid groups + assertTrue(testSelector.checkQueryUserGroups("user2", queryGroups)); + } + + @Test + public void testStarInNegativeUsers() throws Exception { + usersValue.add("*:-"); + usersValue.add("user1:+"); + + List<String> expectedAllowedUsers = new ArrayList<>(); + expectedAllowedUsers.add("user1"); + + List<String> expectedDisAllowedUsers = new ArrayList<>(); + expectedDisAllowedUsers.add("*"); + + aclConfigValue.put("users", usersValue); + AclSelector testSelector = (AclSelector)testCommonHelper(expectedAllowedUsers, emptyList, + expectedDisAllowedUsers, emptyList); + + Set<String> queryGroups = new HashSet<>(); + // Other user with Invalid groups + assertFalse(testSelector.checkQueryUserGroups("user2", queryGroups)); + + // +ve user with invalid groups + assertTrue(testSelector.checkQueryUserGroups("user1", queryGroups)); + } +} diff --git a/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/config/selectors/TestComplexSelectors.java b/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/config/selectors/TestComplexSelectors.java new file mode 100644 index 000000000..d381d05a0 --- /dev/null +++ b/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/config/selectors/TestComplexSelectors.java @@ -0,0 +1,247 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.drill.exec.resourcemgr.config.selectors; + +import com.typesafe.config.Config; +import com.typesafe.config.ConfigFactory; +import com.typesafe.config.ConfigValueFactory; +import org.apache.drill.categories.ResourceManagerTest; +import org.apache.drill.exec.ExecConstants; +import org.apache.drill.exec.ops.QueryContext; +import org.apache.drill.exec.resourcemgr.config.exception.RMConfigException; +import org.apache.drill.exec.server.options.OptionValue; +import org.junit.After; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +@Category(ResourceManagerTest.class) +public final class TestComplexSelectors { + + private static final Map<String, String> tagSelectorConfig1 = new HashMap<>(); + + private static final Map<String, String> tagSelectorConfig2 = new HashMap<>(); + + private static final QueryContext mockContext = mock(QueryContext.class); + + private static final List<Object> complexSelectorValue = new ArrayList<>(); + + @BeforeClass + public static void classLevelSetup() { + tagSelectorConfig1.put("tag", "small"); + tagSelectorConfig2.put("tag", "large"); + } + + @After + public void testLevelSetup() { + complexSelectorValue.clear(); + } + + private ResourcePoolSelector testOrCommonHelper() throws RMConfigException { + Config testConfig = ConfigFactory.empty().withValue("or", ConfigValueFactory.fromIterable(complexSelectorValue)); + final ResourcePoolSelector testSelector = ResourcePoolSelectorFactory.createSelector(testConfig); + assertTrue("TestSelector is not a OrSelector", testSelector instanceof OrSelector); + return testSelector; + } + + private ResourcePoolSelector testAndCommonHelper() throws RMConfigException { + Config testConfig = ConfigFactory.empty().withValue("and", ConfigValueFactory.fromIterable(complexSelectorValue)); + final ResourcePoolSelector testSelector = ResourcePoolSelectorFactory.createSelector(testConfig); + assertTrue("TestSelector is not a AndSelector", testSelector instanceof AndSelector); + return testSelector; + } + + @Test(expected = RMConfigException.class) + public void testOrSelectorSingleValidSelector() throws Exception { + // setup complexSelectorValue + complexSelectorValue.add(tagSelectorConfig1); + testOrCommonHelper(); + } + + @Test(expected = RMConfigException.class) + public void testOrSelectorEmptyValue() throws Exception { + testOrCommonHelper(); + } + + @Test(expected = RMConfigException.class) + public void testAndSelectorEmptyValue() throws Exception { + testAndCommonHelper(); + } + + @Test(expected = RMConfigException.class) + public void testOrSelectorStringValue() throws Exception { + Config testConfig = ConfigFactory.empty().withValue("or", ConfigValueFactory.fromAnyRef("dummy")); + ResourcePoolSelectorFactory.createSelector(testConfig); + } + + @Test(expected = RMConfigException.class) + public void testAndSelectorStringValue() throws Exception { + Config testConfig = ConfigFactory.empty().withValue("and", ConfigValueFactory.fromAnyRef("dummy")); + ResourcePoolSelectorFactory.createSelector(testConfig); + } + + @Test + public void testOrSelectorWithTwoValidSelector() throws Exception { + // setup complexSelectorValue + complexSelectorValue.add(tagSelectorConfig1); + complexSelectorValue.add(tagSelectorConfig2); + + // Test for OrSelector + ResourcePoolSelector testSelector = testOrCommonHelper(); + + // Test successful selection with OrSelector + OptionValue testOption = OptionValue.create(OptionValue.AccessibleScopes.SESSION_AND_QUERY, + ExecConstants.RM_QUERY_TAGS_KEY, "small", OptionValue.OptionScope.SESSION); + when(mockContext.getOption(ExecConstants.RM_QUERY_TAGS_KEY)).thenReturn(testOption); + assertTrue(testSelector.isQuerySelected(mockContext)); + + // Test unsuccessful selection with OrSelector + testOption = OptionValue.create(OptionValue.AccessibleScopes.SESSION_AND_QUERY, + ExecConstants.RM_QUERY_TAGS_KEY, "medium", OptionValue.OptionScope.SESSION); + when(mockContext.getOption(ExecConstants.RM_QUERY_TAGS_KEY)).thenReturn(testOption); + assertFalse(testSelector.isQuerySelected(mockContext)); + } + + @Test + public void testAndSelectorWithTwoValidSelector() throws Exception { + // setup complexSelectorValue + complexSelectorValue.add(tagSelectorConfig1); + complexSelectorValue.add(tagSelectorConfig2); + + // Test for AndSelector + ResourcePoolSelector testSelector = testAndCommonHelper(); + + // Test successful selection with AndSelector + OptionValue testOption = OptionValue.create(OptionValue.AccessibleScopes.SESSION_AND_QUERY, + ExecConstants.RM_QUERY_TAGS_KEY, "small,large", OptionValue.OptionScope.SESSION); + when(mockContext.getOption(ExecConstants.RM_QUERY_TAGS_KEY)).thenReturn(testOption); + assertTrue(testSelector.isQuerySelected(mockContext)); + + // Test unsuccessful selection with AndSelector + testOption = OptionValue.create(OptionValue.AccessibleScopes.SESSION_AND_QUERY, + ExecConstants.RM_QUERY_TAGS_KEY, "small", OptionValue.OptionScope.SESSION); + when(mockContext.getOption(ExecConstants.RM_QUERY_TAGS_KEY)).thenReturn(testOption); + assertFalse(testSelector.isQuerySelected(mockContext)); + } + + @Test + public void testAndSelectorWithNotEqualSelector() throws Exception { + // setup NotEqualSelector config + Map<String, Object> notEqualConfig = new HashMap<>(); + notEqualConfig.put("not_equal", tagSelectorConfig1); + + // setup complex selector value + complexSelectorValue.add(notEqualConfig); + complexSelectorValue.add(tagSelectorConfig2); + + // Test for AndSelector + ResourcePoolSelector testSelector = testAndCommonHelper(); + + // Test successful selection with AndSelector + OptionValue testOption = OptionValue.create(OptionValue.AccessibleScopes.SESSION_AND_QUERY, + ExecConstants.RM_QUERY_TAGS_KEY, "large", OptionValue.OptionScope.SESSION); + when(mockContext.getOption(ExecConstants.RM_QUERY_TAGS_KEY)).thenReturn(testOption); + assertTrue(testSelector.isQuerySelected(mockContext)); + + // Test unsuccessful selection with AndSelector + testOption = OptionValue.create(OptionValue.AccessibleScopes.SESSION_AND_QUERY, + ExecConstants.RM_QUERY_TAGS_KEY, "small", OptionValue.OptionScope.SESSION); + when(mockContext.getOption(ExecConstants.RM_QUERY_TAGS_KEY)).thenReturn(testOption); + assertFalse(testSelector.isQuerySelected(mockContext)); + } + + @Test + public void testORSelectorWithNotEqualSelector() throws Exception { + // setup NotEqualSelector config + Map<String, Object> notEqualConfig = new HashMap<>(); + notEqualConfig.put("not_equal", tagSelectorConfig1); + + // setup complex selector value + complexSelectorValue.add(notEqualConfig); + complexSelectorValue.add(tagSelectorConfig2); + + // Test for AndSelector + ResourcePoolSelector testSelector = testOrCommonHelper(); + + // Test successful selection with AndSelector + OptionValue testOption = OptionValue.create(OptionValue.AccessibleScopes.SESSION_AND_QUERY, + ExecConstants.RM_QUERY_TAGS_KEY, "medium", OptionValue.OptionScope.SESSION); + when(mockContext.getOption(ExecConstants.RM_QUERY_TAGS_KEY)).thenReturn(testOption); + assertTrue(testSelector.isQuerySelected(mockContext)); + + testOption = OptionValue.create(OptionValue.AccessibleScopes.SESSION_AND_QUERY, + ExecConstants.RM_QUERY_TAGS_KEY, "large", OptionValue.OptionScope.SESSION); + when(mockContext.getOption(ExecConstants.RM_QUERY_TAGS_KEY)).thenReturn(testOption); + assertTrue(testSelector.isQuerySelected(mockContext)); + + // Test unsuccessful selection with AndSelector + testOption = OptionValue.create(OptionValue.AccessibleScopes.SESSION_AND_QUERY, + ExecConstants.RM_QUERY_TAGS_KEY, "small", OptionValue.OptionScope.SESSION); + when(mockContext.getOption(ExecConstants.RM_QUERY_TAGS_KEY)).thenReturn(testOption); + assertFalse(testSelector.isQuerySelected(mockContext)); + } + + @Test + public void testAndSelectorWithOrSelector() throws Exception { + // setup OrSelector config + List<Object> orConfigValue = new ArrayList<>(); + orConfigValue.add(tagSelectorConfig1); + orConfigValue.add(tagSelectorConfig2); + + Map<String, Object> orConfig = new HashMap<>(); + orConfig.put("or", orConfigValue); + + // get another TagSelector config + Map<String, String> tagSelectorConfig3 = new HashMap<>(); + tagSelectorConfig3.put("tag", "medium"); + + // setup complex selector value + complexSelectorValue.add(orConfig); + complexSelectorValue.add(tagSelectorConfig3); + + // Test for AndSelector + ResourcePoolSelector testSelector = testAndCommonHelper(); + + // Test successful selection with AndSelector + OptionValue testOption = OptionValue.create(OptionValue.AccessibleScopes.SESSION_AND_QUERY, + ExecConstants.RM_QUERY_TAGS_KEY, "small,medium", OptionValue.OptionScope.SESSION); + when(mockContext.getOption(ExecConstants.RM_QUERY_TAGS_KEY)).thenReturn(testOption); + assertTrue(testSelector.isQuerySelected(mockContext)); + + testOption = OptionValue.create(OptionValue.AccessibleScopes.SESSION_AND_QUERY, + ExecConstants.RM_QUERY_TAGS_KEY, "large,medium", OptionValue.OptionScope.SESSION); + when(mockContext.getOption(ExecConstants.RM_QUERY_TAGS_KEY)).thenReturn(testOption); + assertTrue(testSelector.isQuerySelected(mockContext)); + + // Test unsuccessful selection with AndSelector + testOption = OptionValue.create(OptionValue.AccessibleScopes.SESSION_AND_QUERY, + ExecConstants.RM_QUERY_TAGS_KEY, "small,verylarge", OptionValue.OptionScope.SESSION); + when(mockContext.getOption(ExecConstants.RM_QUERY_TAGS_KEY)).thenReturn(testOption); + assertFalse(testSelector.isQuerySelected(mockContext)); + } +} diff --git a/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/config/selectors/TestNotEqualSelector.java b/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/config/selectors/TestNotEqualSelector.java new file mode 100644 index 000000000..b3053aa22 --- /dev/null +++ b/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/config/selectors/TestNotEqualSelector.java @@ -0,0 +1,120 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.drill.exec.resourcemgr.config.selectors; + +import com.typesafe.config.Config; +import com.typesafe.config.ConfigFactory; +import com.typesafe.config.ConfigValueFactory; +import org.apache.drill.categories.ResourceManagerTest; +import org.apache.drill.exec.ExecConstants; +import org.apache.drill.exec.ops.QueryContext; +import org.apache.drill.exec.resourcemgr.config.exception.RMConfigException; +import org.apache.drill.exec.server.options.OptionValue; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +@Category(ResourceManagerTest.class) +public final class TestNotEqualSelector { + + private ResourcePoolSelector testCommonHelper(Map<String, ? extends Object> selectorValue) throws RMConfigException { + Config testConfig = ConfigFactory.empty() + .withValue("not_equal", ConfigValueFactory.fromMap(selectorValue)); + final ResourcePoolSelector testSelector = ResourcePoolSelectorFactory.createSelector(testConfig); + assertTrue("TestSelector is not a not_equal selector", testSelector instanceof NotEqualSelector); + return testSelector; + } + + @Test + public void testValidNotEqualTagSelector() throws Exception { + Map<String, String> tagSelectorConfig = new HashMap<>(); + tagSelectorConfig.put("tag", "marketing"); + NotEqualSelector testSelector = (NotEqualSelector)testCommonHelper(tagSelectorConfig); + assertTrue("Expected child selector type to be TagSelector", + testSelector.getPoolSelector() instanceof TagSelector); + + QueryContext mockContext = mock(QueryContext.class); + + // Test successful selection + OptionValue testOption = OptionValue.create(OptionValue.AccessibleScopes.SESSION_AND_QUERY, ExecConstants + .RM_QUERY_TAGS_KEY, "small,large", OptionValue.OptionScope.SESSION); + when(mockContext.getOption(ExecConstants.RM_QUERY_TAGS_KEY)).thenReturn(testOption); + assertTrue(testSelector.isQuerySelected(mockContext)); + + // Test unsuccessful selection + testOption = OptionValue.create(OptionValue.AccessibleScopes.SESSION_AND_QUERY, ExecConstants + .RM_QUERY_TAGS_KEY, "marketing", OptionValue.OptionScope.SESSION); + when(mockContext.getOption(ExecConstants.RM_QUERY_TAGS_KEY)).thenReturn(testOption); + assertFalse(testSelector.isQuerySelected(mockContext)); + } + + @Test + public void testValidNotEqualAclSelector() throws Exception { + Map<String, List<String>> aclSelectorValue = new HashMap<>(); + List<String> usersValue = new ArrayList<>(); + usersValue.add("user1"); + usersValue.add("user2"); + + List<String> groupsValue = new ArrayList<>(); + groupsValue.add("group1"); + groupsValue.add("group2"); + + aclSelectorValue.put("users", usersValue); + aclSelectorValue.put("groups", groupsValue); + + Map<String, Map<String, List<String>>> aclSelectorConfig = new HashMap<>(); + aclSelectorConfig.put("acl", aclSelectorValue); + NotEqualSelector testSelector = (NotEqualSelector)testCommonHelper(aclSelectorConfig); + assertTrue("Expected child selector type to be TagSelector", + testSelector.getPoolSelector() instanceof AclSelector); + } + + @Test(expected = RMConfigException.class) + public void testInValidNotEqualAclSelector() throws Exception { + Map<String, List<String>> aclSelectorValue = new HashMap<>(); + + aclSelectorValue.put("users", new ArrayList<>()); + aclSelectorValue.put("groups", new ArrayList<>()); + + Map<String, Map<String, List<String>>> aclSelectorConfig = new HashMap<>(); + aclSelectorConfig.put("acl", aclSelectorValue); + testCommonHelper(aclSelectorConfig); + } + + @Test(expected = RMConfigException.class) + public void testNotEqualSelectorEmptyValue() throws Exception { + Map<String, String> notEqualSelectorValue = new HashMap<>(); + testCommonHelper(notEqualSelectorValue); + } + + @Test(expected = RMConfigException.class) + public void testNotEqualSelectorStringValue() throws Exception { + Config testConfig = ConfigFactory.empty() + .withValue("not_equal", ConfigValueFactory.fromAnyRef("null")); + ResourcePoolSelectorFactory.createSelector(testConfig); + } +} diff --git a/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/config/selectors/TestResourcePoolSelectors.java b/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/config/selectors/TestResourcePoolSelectors.java new file mode 100644 index 000000000..0a2d121c9 --- /dev/null +++ b/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/config/selectors/TestResourcePoolSelectors.java @@ -0,0 +1,46 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.drill.exec.resourcemgr.config.selectors; + +import com.typesafe.config.Config; +import com.typesafe.config.ConfigFactory; +import org.apache.drill.categories.ResourceManagerTest; +import org.apache.drill.exec.resourcemgr.config.exception.RMConfigException; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import static org.junit.Assert.assertTrue; + +@Category(ResourceManagerTest.class) +public class TestResourcePoolSelectors { + + @Test + public void testNullSelectorConfig() throws Exception { + final ResourcePoolSelector testSelector = ResourcePoolSelectorFactory.createSelector(null); + assertTrue("TestSelector with null config is not of type Default Selector", + testSelector instanceof DefaultSelector); + assertTrue("DefaultSelector type is not default", + testSelector.getSelectorType() == ResourcePoolSelector.SelectorType.DEFAULT); + } + + @Test(expected = RMConfigException.class) + public void testEmptySelectorConfig() throws Exception { + Config testConfig = ConfigFactory.empty(); + ResourcePoolSelectorFactory.createSelector(testConfig); + } +} diff --git a/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/config/selectors/TestTagSelector.java b/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/config/selectors/TestTagSelector.java new file mode 100644 index 000000000..d93731513 --- /dev/null +++ b/exec/java-exec/src/test/java/org/apache/drill/exec/resourcemgr/config/selectors/TestTagSelector.java @@ -0,0 +1,84 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.drill.exec.resourcemgr.config.selectors; + +import com.typesafe.config.Config; +import com.typesafe.config.ConfigFactory; +import com.typesafe.config.ConfigValueFactory; +import org.apache.drill.categories.ResourceManagerTest; +import org.apache.drill.exec.ExecConstants; +import org.apache.drill.exec.ops.QueryContext; +import org.apache.drill.exec.resourcemgr.config.exception.RMConfigException; +import org.apache.drill.exec.server.options.OptionValue; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +@Category(ResourceManagerTest.class) +public final class TestTagSelector { + + private ResourcePoolSelector testCommonHelper(Object tagValue) throws RMConfigException { + Config testConfig = ConfigFactory.empty() + .withValue("tag", ConfigValueFactory.fromAnyRef(tagValue)); + final ResourcePoolSelector testSelector = ResourcePoolSelectorFactory.createSelector(testConfig); + assertTrue("TestSelector is not a tag selector", testSelector instanceof TagSelector); + assertEquals("Unexpected value for tag selector", tagValue, ((TagSelector) testSelector).getTagValue()); + return testSelector; + } + + // Tests for TagSelector + @Test + public void testValidTagSelector() throws Exception { + testCommonHelper("marketing"); + } + + @Test(expected = RMConfigException.class) + public void testInValidTagSelector() throws Exception { + testCommonHelper(""); + testCommonHelper(null); + } + + @Test + public void testTagSelectorWithQueryTags() throws Exception { + QueryContext mockContext = mock(QueryContext.class); + TagSelector testSelector = (TagSelector) testCommonHelper("small"); + + // Test successful selection + OptionValue testOption = OptionValue.create(OptionValue.AccessibleScopes.SESSION_AND_QUERY, ExecConstants + .RM_QUERY_TAGS_KEY, "small,large", OptionValue.OptionScope.SESSION); + when(mockContext.getOption(ExecConstants.RM_QUERY_TAGS_KEY)).thenReturn(testOption); + assertTrue(testSelector.isQuerySelected(mockContext)); + + // Test empty query tags + testOption = OptionValue.create(OptionValue.AccessibleScopes.SESSION_AND_QUERY, ExecConstants + .RM_QUERY_TAGS_KEY, "", OptionValue.OptionScope.SESSION); + when(mockContext.getOption(ExecConstants.RM_QUERY_TAGS_KEY)).thenReturn(testOption); + assertFalse(testSelector.isQuerySelected(mockContext)); + + // Test different query tags + testOption = OptionValue.create(OptionValue.AccessibleScopes.SESSION_AND_QUERY, ExecConstants + .RM_QUERY_TAGS_KEY, "medium", OptionValue.OptionScope.SESSION); + when(mockContext.getOption(ExecConstants.RM_QUERY_TAGS_KEY)).thenReturn(testOption); + assertFalse(testSelector.isQuerySelected(mockContext)); + } +} |