use of org.apache.drill.exec.resourcemgr.config.QueueAssignmentResult in project drill by apache.
the class TestResourcePoolTree method testTreeWithLeafAndIntermediatePool.
@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();
}
Aggregations