/**
 * 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.hadoop.yarn.server.resourcemanager.monitor.capacity;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
import org.apache.hadoop.yarn.api.records.ApplicationId;
import org.apache.hadoop.yarn.api.records.Container;
import org.apache.hadoop.yarn.api.records.ContainerId;
import org.apache.hadoop.yarn.api.records.NodeId;
import org.apache.hadoop.yarn.api.records.Priority;
import org.apache.hadoop.yarn.api.records.Resource;
import org.apache.hadoop.yarn.event.Dispatcher;
import org.apache.hadoop.yarn.event.EventHandler;
import org.apache.hadoop.yarn.server.resourcemanager.RMContext;
import org.apache.hadoop.yarn.server.resourcemanager.nodelabels.RMNodeLabelsManager;
import org.apache.hadoop.yarn.server.resourcemanager.rmcontainer.RMContainer;
import org.apache.hadoop.yarn.server.resourcemanager.rmcontainer.RMContainerImpl;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.ResourceUsage;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.SchedulerNode;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.policy.QueueOrderingPolicy;
import org.apache.hadoop.yarn.server.scheduler.SchedulerRequestKey;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.CSQueue;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.CapacityScheduler;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.CapacitySchedulerConfiguration;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.LeafQueue;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.ParentQueue;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.QueueCapacities;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.SchedulingMode;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.UsersManager.User;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.preemption.PreemptionManager;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.common.fica.FiCaSchedulerApp;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.common.fica.FiCaSchedulerNode;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.event.ContainerPreemptEvent;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.policy.OrderingPolicy;
import org.apache.hadoop.yarn.util.Clock;
import org.apache.hadoop.yarn.util.resource.DefaultResourceCalculator;
import org.apache.hadoop.yarn.util.resource.DominantResourceCalculator;
import org.apache.hadoop.yarn.util.resource.ResourceCalculator;
import org.apache.hadoop.yarn.util.resource.Resources;
import org.junit.Assert;
import org.junit.Before;
import org.mockito.ArgumentMatcher;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.TreeSet;
import java.util.concurrent.locks.ReentrantReadWriteLock;
import org.apache.hadoop.yarn.event.Event;

import static org.mockito.Matchers.any;
import static org.mockito.Matchers.anyString;
import static org.mockito.Matchers.eq;
import static org.mockito.Matchers.isA;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

public class ProportionalCapacityPreemptionPolicyMockFramework {
  static final Log LOG =
      LogFactory.getLog(TestProportionalCapacityPreemptionPolicyForNodePartitions.class);
  final String ROOT = CapacitySchedulerConfiguration.ROOT;

  Map<String, CSQueue> nameToCSQueues = null;
  Map<String, Resource> partitionToResource = null;
  Map<NodeId, FiCaSchedulerNode> nodeIdToSchedulerNodes = null;
  RMNodeLabelsManager nlm = null;
  RMContext rmContext = null;

  ResourceCalculator rc = new DefaultResourceCalculator();
  Clock mClock = null;
  CapacitySchedulerConfiguration conf = null;
  CapacityScheduler cs = null;
  @SuppressWarnings("rawtypes")
  EventHandler<Event> mDisp = null;
  ProportionalCapacityPreemptionPolicy policy = null;
  Resource clusterResource = null;

  @SuppressWarnings("unchecked")
  @Before
  public void setup() {
    org.apache.log4j.Logger.getRootLogger().setLevel(
        org.apache.log4j.Level.DEBUG);

    conf = new CapacitySchedulerConfiguration(new Configuration(false));
    conf.setLong(
        CapacitySchedulerConfiguration.PREEMPTION_WAIT_TIME_BEFORE_KILL, 10000);
    conf.setLong(CapacitySchedulerConfiguration.PREEMPTION_MONITORING_INTERVAL,
        3000);
    // report "ideal" preempt
    conf.setFloat(CapacitySchedulerConfiguration.TOTAL_PREEMPTION_PER_ROUND,
        (float) 1.0);
    conf.setFloat(
        CapacitySchedulerConfiguration.PREEMPTION_NATURAL_TERMINATION_FACTOR,
        (float) 1.0);

    mClock = mock(Clock.class);
    cs = mock(CapacityScheduler.class);
    when(cs.getResourceCalculator()).thenReturn(rc);
    when(cs.getPreemptionManager()).thenReturn(new PreemptionManager());
    when(cs.getConfiguration()).thenReturn(conf);

    nlm = mock(RMNodeLabelsManager.class);
    mDisp = mock(EventHandler.class);

    rmContext = mock(RMContext.class);
    when(rmContext.getNodeLabelManager()).thenReturn(nlm);
    Dispatcher disp = mock(Dispatcher.class);
    when(rmContext.getDispatcher()).thenReturn(disp);
    when(disp.getEventHandler()).thenReturn(mDisp);
    when(cs.getRMContext()).thenReturn(rmContext);

    partitionToResource = new HashMap<>();
    nodeIdToSchedulerNodes = new HashMap<>();
    nameToCSQueues = new HashMap<>();
  }

  public void buildEnv(String labelsConfig, String nodesConfig,
      String queuesConfig, String appsConfig) throws IOException {
    buildEnv(labelsConfig, nodesConfig, queuesConfig, appsConfig, false);
  }

  public void buildEnv(String labelsConfig, String nodesConfig,
      String queuesConfig, String appsConfig,
      boolean useDominantResourceCalculator) throws IOException {
    if (useDominantResourceCalculator) {
      when(cs.getResourceCalculator()).thenReturn(
          new DominantResourceCalculator());
    }
    mockNodeLabelsManager(labelsConfig);
    mockSchedulerNodes(nodesConfig);
    for (NodeId nodeId : nodeIdToSchedulerNodes.keySet()) {
      when(cs.getSchedulerNode(nodeId)).thenReturn(
          nodeIdToSchedulerNodes.get(nodeId));
    }
    List<FiCaSchedulerNode> allNodes = new ArrayList<>(
        nodeIdToSchedulerNodes.values());
    when(cs.getAllNodes()).thenReturn(allNodes);
    ParentQueue root = mockQueueHierarchy(queuesConfig);
    when(cs.getRootQueue()).thenReturn(root);
    when(cs.getClusterResource()).thenReturn(clusterResource);
    mockApplications(appsConfig);

    policy = new ProportionalCapacityPreemptionPolicy(rmContext, cs,
        mClock);
  }

  private void mockContainers(String containersConfig, FiCaSchedulerApp app,
      ApplicationAttemptId attemptId, String queueName,
      List<RMContainer> reservedContainers, List<RMContainer> liveContainers) {
    int containerId = 1;
    int start = containersConfig.indexOf("=") + 1;
    int end = -1;

    Resource used = Resource.newInstance(0, 0);
    Resource pending = Resource.newInstance(0, 0);
    Priority pri = Priority.newInstance(0);

    while (start < containersConfig.length()) {
      while (start < containersConfig.length()
          && containersConfig.charAt(start) != '(') {
        start++;
      }
      if (start >= containersConfig.length()) {
        throw new IllegalArgumentException(
            "Error containers specification, line=" + containersConfig);
      }
      end = start + 1;
      while (end < containersConfig.length()
          && containersConfig.charAt(end) != ')') {
        end++;
      }
      if (end >= containersConfig.length()) {
        throw new IllegalArgumentException(
            "Error containers specification, line=" + containersConfig);
      }

      // now we found start/end, get container values
      String[] values = containersConfig.substring(start + 1, end).split(",");
      if (values.length < 6 || values.length > 8) {
        throw new IllegalArgumentException("Format to define container is:"
            + "(priority,resource,host,expression,repeat,reserved, pending)");
      }
      pri.setPriority(Integer.valueOf(values[0]));
      Resource res = parseResourceFromString(values[1]);
      NodeId host = NodeId.newInstance(values[2], 1);
      String label = values[3];
      String userName = "user";
      int repeat = Integer.valueOf(values[4]);
      boolean reserved = Boolean.valueOf(values[5]);
      if (values.length >= 7) {
        Resources.addTo(pending, parseResourceFromString(values[6]));
      }
      if (values.length == 8) {
        userName = values[7];
      }

      for (int i = 0; i < repeat; i++) {
        Container c = mock(Container.class);
        Resources.addTo(used, res);
        when(c.getResource()).thenReturn(res);
        when(c.getPriority()).thenReturn(pri);
        SchedulerRequestKey sk = SchedulerRequestKey.extractFrom(c);
        RMContainerImpl rmc = mock(RMContainerImpl.class);
        when(rmc.getAllocatedSchedulerKey()).thenReturn(sk);
        when(rmc.getAllocatedNode()).thenReturn(host);
        when(rmc.getNodeLabelExpression()).thenReturn(label);
        when(rmc.getAllocatedResource()).thenReturn(res);
        when(rmc.getContainer()).thenReturn(c);
        when(rmc.getApplicationAttemptId()).thenReturn(attemptId);
        when(rmc.getQueueName()).thenReturn(queueName);
        final ContainerId cId = ContainerId.newContainerId(attemptId,
            containerId);
        when(rmc.getContainerId()).thenReturn(cId);
        doAnswer(new Answer<Integer>() {
          @Override
          public Integer answer(InvocationOnMock invocation) throws Throwable {
            return cId.compareTo(
                ((RMContainer) invocation.getArguments()[0]).getContainerId());
          }
        }).when(rmc).compareTo(any(RMContainer.class));

        if (containerId == 1) {
          when(rmc.isAMContainer()).thenReturn(true);
          when(app.getAMResource(label)).thenReturn(res);
          when(app.getAppAMNodePartitionName()).thenReturn(label);
        }

        if (reserved) {
          reservedContainers.add(rmc);
          when(rmc.getReservedResource()).thenReturn(res);
        } else {
          liveContainers.add(rmc);
        }

        // Add container to scheduler-node
        addContainerToSchedulerNode(host, rmc, reserved);

        // If this is a non-exclusive allocation
        String partition = null;
        if (label.isEmpty()
            && !(partition = nodeIdToSchedulerNodes.get(host).getPartition())
                .isEmpty()) {
          LeafQueue queue = (LeafQueue) nameToCSQueues.get(queueName);
          Map<String, TreeSet<RMContainer>> ignoreExclusivityContainers = queue
              .getIgnoreExclusivityRMContainers();
          if (!ignoreExclusivityContainers.containsKey(partition)) {
            ignoreExclusivityContainers.put(partition,
                new TreeSet<RMContainer>());
          }
          ignoreExclusivityContainers.get(partition).add(rmc);
        }
        LOG.debug("add container to app=" + attemptId + " res=" + res + " node="
            + host + " nodeLabelExpression=" + label + " partition="
            + partition);

        containerId++;
      }

      // If app has 0 container, and it has only pending, still make sure to
      // update label.
      if (repeat == 0) {
        when(app.getAppAMNodePartitionName()).thenReturn(label);
      }

      // Some more app specific aggregated data can be better filled here.
      when(app.getPriority()).thenReturn(pri);
      when(app.getUser()).thenReturn(userName);
      when(app.getCurrentConsumption()).thenReturn(used);
      when(app.getCurrentReservation())
          .thenReturn(Resources.createResource(0, 0));

      Map<String, Resource> pendingForDefaultPartition =
          new HashMap<String, Resource>();
      // Add for default partition for now.
      pendingForDefaultPartition.put(label, pending);
      when(app.getTotalPendingRequestsPerPartition())
          .thenReturn(pendingForDefaultPartition);

      // need to set pending resource in resource usage as well
      ResourceUsage ru = new ResourceUsage();
      ru.setUsed(label, used);
      when(app.getAppAttemptResourceUsage()).thenReturn(ru);

      start = end + 1;
    }
  }

  /**
   * Format is:
   * <pre>
   * queueName\t  // app1
   * (priority,resource,host,expression,#repeat,reserved)
   * (priority,resource,host,expression,#repeat,reserved);
   * queueName\t  // app2
   * </pre>
   */
  private void mockApplications(String appsConfig) {
    int id = 1;
    HashMap<String, HashSet<String>> userMap = new HashMap<String, HashSet<String>>();
    HashMap<String, HashMap<String, HashMap<String, ResourceUsage>>> userResourceUsagePerLabel = new HashMap<>();
    LeafQueue queue = null;
    int mulp = -1;
    for (String a : appsConfig.split(";")) {
      String[] strs = a.split("\t");
      String queueName = strs[0];
      if (mulp <= 0 && strs.length > 2 && strs[2] != null) {
        mulp = 100 / (new Integer(strs[2]).intValue());
      }

      // get containers
      List<RMContainer> liveContainers = new ArrayList<RMContainer>();
      List<RMContainer> reservedContainers = new ArrayList<RMContainer>();
      ApplicationId appId = ApplicationId.newInstance(0L, id);
      ApplicationAttemptId appAttemptId = ApplicationAttemptId
          .newInstance(appId, 1);

      FiCaSchedulerApp app = mock(FiCaSchedulerApp.class);
      when(app.getAMResource(anyString()))
          .thenReturn(Resources.createResource(0, 0));
      mockContainers(strs[1], app, appAttemptId, queueName, reservedContainers,
          liveContainers);
      LOG.debug("Application mock: queue: " + queueName + ", appId:" + appId);

      when(app.getLiveContainers()).thenReturn(liveContainers);
      when(app.getReservedContainers()).thenReturn(reservedContainers);
      when(app.getApplicationAttemptId()).thenReturn(appAttemptId);
      when(app.getApplicationId()).thenReturn(appId);
      when(app.getQueueName()).thenReturn(queueName);

      // add to LeafQueue
      queue = (LeafQueue) nameToCSQueues.get(queueName);
      queue.getApplications().add(app);
      queue.getAllApplications().add(app);
      when(queue.getMinimumAllocation())
      .thenReturn(Resource.newInstance(1,1));
      when(app.getCSLeafQueue()).thenReturn(queue);

      HashSet<String> users = userMap.get(queueName);
      if (null == users) {
        users = new HashSet<String>();
        userMap.put(queueName, users);
      }
      users.add(app.getUser());

      String label = app.getAppAMNodePartitionName();

      // Get label to queue
      HashMap<String, HashMap<String, ResourceUsage>> userResourceUsagePerQueue = userResourceUsagePerLabel
          .get(label);
      if (null == userResourceUsagePerQueue) {
        userResourceUsagePerQueue = new HashMap<>();
        userResourceUsagePerLabel.put(label, userResourceUsagePerQueue);
      }

      // Get queue to user based resource map
      HashMap<String, ResourceUsage> userResourceUsage = userResourceUsagePerQueue
          .get(queueName);
      if (null == userResourceUsage) {
        userResourceUsage = new HashMap<>();
        userResourceUsagePerQueue.put(queueName, userResourceUsage);
      }

      // Get user to its resource usage.
      ResourceUsage usage = userResourceUsage.get(app.getUser());
      if (null == usage) {
        usage = new ResourceUsage();
        userResourceUsage.put(app.getUser(), usage);
      }

      usage.incAMUsed(app.getAMResource(label));
      usage.incUsed(app.getAppAttemptResourceUsage().getUsed(label));
      id++;
    }

    for (String label : userResourceUsagePerLabel.keySet()) {
      for (String queueName : userMap.keySet()) {
        queue = (LeafQueue) nameToCSQueues.get(queueName);
        // Currently we have user-limit test support only for default label.
        Resource totResoucePerPartition = partitionToResource.get("");
        Resource capacity = Resources.multiply(totResoucePerPartition,
            queue.getQueueCapacities().getAbsoluteCapacity());
        HashSet<String> users = userMap.get(queue.getQueueName());
        when(queue.getAllUsers()).thenReturn(users);
        Resource userLimit;
        if (mulp > 0) {
          userLimit = Resources.divideAndCeil(rc, capacity, mulp);
        } else {
          userLimit = Resources.divideAndCeil(rc, capacity,
              users.size());
        }
        LOG.debug("Updating user-limit from mock: totResoucePerPartition="
            + totResoucePerPartition + ", capacity=" + capacity
            + ", users.size()=" + users.size() + ", userlimit= " + userLimit
            + ",label= " + label + ",queueName= " + queueName);

        HashMap<String, ResourceUsage> userResourceUsage = userResourceUsagePerLabel
            .get(label).get(queueName);
        for (String userName : users) {
          User user = new User(userName);
          if (userResourceUsage != null) {
            user.setResourceUsage(userResourceUsage.get(userName));
          }
          when(queue.getUser(eq(userName))).thenReturn(user);
          when(queue.getResourceLimitForAllUsers(eq(userName),
              any(Resource.class), anyString(), any(SchedulingMode.class)))
                  .thenReturn(userLimit);
        }
      }
    }
  }

  private void addContainerToSchedulerNode(NodeId nodeId, RMContainer container,
      boolean isReserved) {
    SchedulerNode node = nodeIdToSchedulerNodes.get(nodeId);
    assert node != null;

    if (isReserved) {
      when(node.getReservedContainer()).thenReturn(container);
    } else {
      node.getCopiedListOfRunningContainers().add(container);
      Resources.subtractFrom(node.getUnallocatedResource(),
          container.getAllocatedResource());
    }
  }

  /**
   * Format is:
   * host1=partition[ res=resource];
   * host2=partition[ res=resource];
   */
  private void mockSchedulerNodes(String schedulerNodesConfigStr)
      throws IOException {
    String[] nodesConfigStrArray = schedulerNodesConfigStr.split(";");
    for (String p : nodesConfigStrArray) {
      String[] arr = p.split(" ");

      NodeId nodeId = NodeId.newInstance(arr[0].substring(0, arr[0].indexOf("=")), 1);
      String partition = arr[0].substring(arr[0].indexOf("=") + 1, arr[0].length());

      FiCaSchedulerNode sn = mock(FiCaSchedulerNode.class);
      when(sn.getNodeID()).thenReturn(nodeId);
      when(sn.getPartition()).thenReturn(partition);

      Resource totalRes = Resources.createResource(0);
      if (arr.length > 1) {
        String res = arr[1];
        if (res.contains("res=")) {
          String resSring = res.substring(
              res.indexOf("res=") + "res=".length());
          totalRes = parseResourceFromString(resSring);
        }
      }
      when(sn.getTotalResource()).thenReturn(totalRes);
      when(sn.getUnallocatedResource()).thenReturn(Resources.clone(totalRes));

      // TODO, add settings of killable resources when necessary
      when(sn.getTotalKillableResources()).thenReturn(Resources.none());

      List<RMContainer> liveContainers = new ArrayList<>();
      when(sn.getCopiedListOfRunningContainers()).thenReturn(liveContainers);

      nodeIdToSchedulerNodes.put(nodeId, sn);

      LOG.debug("add scheduler node, id=" + nodeId + ", partition=" + partition);
    }
  }

  /**
   * Format is:
   * <pre>
   * partition0=total_resource,exclusivity;
   * partition1=total_resource,exclusivity;
   * ...
   * </pre>
   */
  private void mockNodeLabelsManager(String nodeLabelsConfigStr) throws IOException {
    String[] partitionConfigArr = nodeLabelsConfigStr.split(";");
    clusterResource = Resources.createResource(0);
    for (String p : partitionConfigArr) {
      String partitionName = p.substring(0, p.indexOf("="));
      Resource res = parseResourceFromString(p.substring(p.indexOf("=") + 1,
          p.indexOf(",")));
      boolean exclusivity =
          Boolean.valueOf(p.substring(p.indexOf(",") + 1, p.length()));
      when(nlm.getResourceByLabel(eq(partitionName), any(Resource.class)))
          .thenReturn(res);
      when(nlm.isExclusiveNodeLabel(eq(partitionName))).thenReturn(exclusivity);

      // add to partition to resource
      partitionToResource.put(partitionName, res);
      LOG.debug("add partition=" + partitionName + " totalRes=" + res
          + " exclusivity=" + exclusivity);
      Resources.addTo(clusterResource, res);
    }

    when(nlm.getClusterNodeLabelNames()).thenReturn(
        partitionToResource.keySet());
  }

  private Resource parseResourceFromString(String p) {
    String[] resource = p.split(":");
    Resource res;
    if (resource.length == 1) {
      res = Resources.createResource(Integer.valueOf(resource[0]));
    } else {
      res = Resources.createResource(Integer.valueOf(resource[0]),
          Integer.valueOf(resource[1]));
    }
    return res;
  }

  /**
   * Format is:
   * <pre>
   * root (<partition-name-1>=[guaranteed max used pending (reserved)],<partition-name-2>=..);
   * -A(...);
   * --A1(...);
   * --A2(...);
   * -B...
   * </pre>
   * ";" splits queues, and there should no empty lines, no extra spaces
   *
   * For each queue, it has configurations to specify capacities (to each
   * partition), format is:
   * <pre>
   * -<queueName> (<labelName1>=[guaranteed max used pending], \
   *               <labelName2>=[guaranteed max used pending])
   *              {key1=value1,key2=value2};  // Additional configs
   * </pre>
   */
  @SuppressWarnings({ "unchecked", "rawtypes" })
  private ParentQueue mockQueueHierarchy(String queueExprs) {
    String[] queueExprArray = queueExprs.split(";");
    ParentQueue rootQueue = null;
    for (int idx = 0; idx < queueExprArray.length; idx++) {
      String q = queueExprArray[idx];
      CSQueue queue;

      // Initialize queue
      if (isParent(queueExprArray, idx)) {
        ParentQueue parentQueue = mock(ParentQueue.class);
        queue = parentQueue;
        List<CSQueue> children = new ArrayList<CSQueue>();
        when(parentQueue.getChildQueues()).thenReturn(children);
        QueueOrderingPolicy policy = mock(QueueOrderingPolicy.class);
        when(policy.getConfigName()).thenReturn(
            CapacitySchedulerConfiguration.QUEUE_PRIORITY_UTILIZATION_ORDERING_POLICY);
        when(parentQueue.getQueueOrderingPolicy()).thenReturn(policy);
      } else {
        LeafQueue leafQueue = mock(LeafQueue.class);
        final TreeSet<FiCaSchedulerApp> apps = new TreeSet<>(
            new Comparator<FiCaSchedulerApp>() {
              @Override
              public int compare(FiCaSchedulerApp a1, FiCaSchedulerApp a2) {
                if (a1.getPriority() != null
                    && !a1.getPriority().equals(a2.getPriority())) {
                  return a1.getPriority().compareTo(a2.getPriority());
                }

                int res = a1.getApplicationId()
                    .compareTo(a2.getApplicationId());
                return res;
              }
            });
        when(leafQueue.getApplications()).thenReturn(apps);
        when(leafQueue.getAllApplications()).thenReturn(apps);
        OrderingPolicy<FiCaSchedulerApp> so = mock(OrderingPolicy.class);
        when(so.getPreemptionIterator()).thenAnswer(new Answer() {
          public Object answer(InvocationOnMock invocation) {
            return apps.descendingIterator();
          }
        });
        when(leafQueue.getOrderingPolicy()).thenReturn(so);

        Map<String, TreeSet<RMContainer>> ignorePartitionContainers =
            new HashMap<>();
        when(leafQueue.getIgnoreExclusivityRMContainers()).thenReturn(
            ignorePartitionContainers);
        queue = leafQueue;
      }

      ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
      when(queue.getReadLock()).thenReturn(lock.readLock());
      setupQueue(queue, q, queueExprArray, idx);
      if (queue.getQueueName().equals(ROOT)) {
        rootQueue = (ParentQueue) queue;
      }
    }
    return rootQueue;
  }

  private void setupQueue(CSQueue queue, String q, String[] queueExprArray,
      int idx) {
    LOG.debug("*** Setup queue, source=" + q);
    String queuePath = null;

    int myLevel = getLevel(q);
    if (0 == myLevel) {
      // It's root
      when(queue.getQueueName()).thenReturn(ROOT);
      queuePath = ROOT;
    }

    String queueName = getQueueName(q);
    when(queue.getQueueName()).thenReturn(queueName);

    // Setup parent queue, and add myself to parentQueue.children-list
    ParentQueue parentQueue = getParentQueue(queueExprArray, idx, myLevel);
    if (null != parentQueue) {
      when(queue.getParent()).thenReturn(parentQueue);
      parentQueue.getChildQueues().add(queue);

      // Setup my path
      queuePath = parentQueue.getQueuePath() + "." + queueName;
    }
    when(queue.getQueuePath()).thenReturn(queuePath);

    QueueCapacities qc = new QueueCapacities(0 == myLevel);
    ResourceUsage ru = new ResourceUsage();

    when(queue.getQueueCapacities()).thenReturn(qc);
    when(queue.getQueueResourceUsage()).thenReturn(ru);

    LOG.debug("Setup queue, name=" + queue.getQueueName() + " path="
        + queue.getQueuePath());
    LOG.debug("Parent=" + (parentQueue == null ? "null" : parentQueue
        .getQueueName()));

    // Setup other fields like used resource, guaranteed resource, etc.
    String capacitySettingStr = q.substring(q.indexOf("(") + 1, q.indexOf(")"));
    for (String s : capacitySettingStr.split(",")) {
      String partitionName = s.substring(0, s.indexOf("="));
      String[] values = s.substring(s.indexOf("[") + 1, s.indexOf("]")).split(" ");
      // Add a small epsilon to capacities to avoid truncate when doing
      // Resources.multiply
      float epsilon = 1e-6f;
      Resource totResoucePerPartition = partitionToResource.get(partitionName);
      float absGuaranteed = Resources.divide(rc, totResoucePerPartition,
          parseResourceFromString(values[0].trim()), totResoucePerPartition)
          + epsilon;
      float absMax = Resources.divide(rc, totResoucePerPartition,
          parseResourceFromString(values[1].trim()), totResoucePerPartition)
          + epsilon;
      float absUsed = Resources.divide(rc, totResoucePerPartition,
          parseResourceFromString(values[2].trim()), totResoucePerPartition)
          + epsilon;
      float used = Resources.divide(rc, totResoucePerPartition,
          parseResourceFromString(values[2].trim()),
          parseResourceFromString(values[0].trim())) + epsilon;
      Resource pending = parseResourceFromString(values[3].trim());
      qc.setAbsoluteCapacity(partitionName, absGuaranteed);
      qc.setAbsoluteMaximumCapacity(partitionName, absMax);
      qc.setAbsoluteUsedCapacity(partitionName, absUsed);
      qc.setUsedCapacity(partitionName, used);
      when(queue.getUsedCapacity()).thenReturn(used);
      ru.setPending(partitionName, pending);
      // Setup reserved resource if it contained by input config
      Resource reserved = Resources.none();
      if(values.length == 5) {
        reserved = parseResourceFromString(values[4].trim());
        ru.setReserved(partitionName, reserved);
      }
      if (!isParent(queueExprArray, idx)) {
        LeafQueue lq = (LeafQueue) queue;
        when(lq.getTotalPendingResourcesConsideringUserLimit(isA(Resource.class),
            isA(String.class), eq(false))).thenReturn(pending);
        when(lq.getTotalPendingResourcesConsideringUserLimit(isA(Resource.class),
            isA(String.class), eq(true))).thenReturn(
            Resources.subtract(pending, reserved));
      }
      ru.setUsed(partitionName, parseResourceFromString(values[2].trim()));

      LOG.debug("Setup queue=" + queueName + " partition=" + partitionName
          + " [abs_guaranteed=" + absGuaranteed + ",abs_max=" + absMax
          + ",abs_used" + absUsed + ",pending_resource=" + pending
          + ", reserved_resource=" + reserved + "]");
    }

    // Setup preemption disabled
    when(queue.getPreemptionDisabled()).thenReturn(
        conf.getPreemptionDisabled(queuePath, false));

    // Setup other queue configurations
    Map<String, String> otherConfigs = getOtherConfigurations(
        queueExprArray[idx]);
    if (otherConfigs.containsKey("priority")) {
      when(queue.getPriority()).thenReturn(
          Priority.newInstance(Integer.valueOf(otherConfigs.get("priority"))));
    } else {
      // set queue's priority to 0 by default
      when(queue.getPriority()).thenReturn(Priority.newInstance(0));
    }

    // Setup disable preemption of queues
    if (otherConfigs.containsKey("disable_preemption")) {
      when(queue.getPreemptionDisabled()).thenReturn(
          Boolean.valueOf(otherConfigs.get("disable_preemption")));
    }

    nameToCSQueues.put(queueName, queue);
    when(cs.getQueue(eq(queueName))).thenReturn(queue);
  }

  /**
   * Get additional queue's configurations
   * @param queueExpr queue expr
   * @return maps of configs
   */
  private Map<String, String> getOtherConfigurations(String queueExpr) {
    if (queueExpr.contains("{")) {
      int left = queueExpr.indexOf('{');
      int right = queueExpr.indexOf('}');

      if (right > left) {
        Map<String, String> configs = new HashMap<>();

        String subStr = queueExpr.substring(left + 1, right);
        for (String kv : subStr.split(",")) {
          if (kv.contains("=")) {
            String key = kv.substring(0, kv.indexOf("="));
            String value = kv.substring(kv.indexOf("=") + 1);
            configs.put(key, value);
          }
        }

        return configs;
      }
    }

    return Collections.EMPTY_MAP;
  }

  /**
   * Level of a queue is how many "-" at beginning, root's level is 0
   */
  private int getLevel(String q) {
    int level = 0; // level = how many "-" at beginning
    while (level < q.length() && q.charAt(level) == '-') {
      level++;
    }
    return level;
  }

  private String getQueueName(String q) {
    int idx = 0;
    // find first != '-' char
    while (idx < q.length() && q.charAt(idx) == '-') {
      idx++;
    }
    if (idx == q.length()) {
      throw new IllegalArgumentException("illegal input:" + q);
    }
    // name = after '-' and before '('
    String name = q.substring(idx, q.indexOf('('));
    if (name.isEmpty()) {
      throw new IllegalArgumentException("queue name shouldn't be empty:" + q);
    }
    if (name.contains(".")) {
      throw new IllegalArgumentException("queue name shouldn't contain '.':"
          + name);
    }
    return name;
  }

  private ParentQueue getParentQueue(String[] queueExprArray, int idx, int myLevel) {
    idx--;
    while (idx >= 0) {
      int level = getLevel(queueExprArray[idx]);
      if (level < myLevel) {
        String parentQueuName = getQueueName(queueExprArray[idx]);
        return (ParentQueue) nameToCSQueues.get(parentQueuName);
      }
      idx--;
    }

    return null;
  }

  /**
   * Get if a queue is ParentQueue
   */
  private boolean isParent(String[] queues, int idx) {
    int myLevel = getLevel(queues[idx]);
    idx++;
    while (idx < queues.length && getLevel(queues[idx]) == myLevel) {
      idx++;
    }
    if (idx >= queues.length || getLevel(queues[idx]) < myLevel) {
      // It's a LeafQueue
      return false;
    } else {
      return true;
    }
  }

  public ApplicationAttemptId getAppAttemptId(int id) {
    ApplicationId appId = ApplicationId.newInstance(0L, id);
    ApplicationAttemptId appAttemptId =
        ApplicationAttemptId.newInstance(appId, 1);
    return appAttemptId;
  }

  public void checkContainerNodesInApp(FiCaSchedulerApp app,
      int expectedContainersNumber, String host) {
    NodeId nodeId = NodeId.newInstance(host, 1);
    int num = 0;
    for (RMContainer c : app.getLiveContainers()) {
      if (c.getAllocatedNode().equals(nodeId)) {
        num++;
      }
    }
    for (RMContainer c : app.getReservedContainers()) {
      if (c.getAllocatedNode().equals(nodeId)) {
        num++;
      }
    }
    Assert.assertEquals(expectedContainersNumber, num);
  }

  public FiCaSchedulerApp getApp(String queueName, int appId) {
    for (FiCaSchedulerApp app : ((LeafQueue) cs.getQueue(queueName))
        .getApplications()) {
      if (app.getApplicationId().getId() == appId) {
        return app;
      }
    }
    return null;
  }

  public void checkAbsCapacities(CSQueue queue, String partition,
      float guaranteed, float max, float used) {
    QueueCapacities qc = queue.getQueueCapacities();
    Assert.assertEquals(guaranteed, qc.getAbsoluteCapacity(partition), 1e-3);
    Assert.assertEquals(max, qc.getAbsoluteMaximumCapacity(partition), 1e-3);
    Assert.assertEquals(used, qc.getAbsoluteUsedCapacity(partition), 1e-3);
  }

  public void checkPendingResource(CSQueue queue, String partition, int pending) {
    ResourceUsage ru = queue.getQueueResourceUsage();
    Assert.assertEquals(pending, ru.getPending(partition).getMemorySize());
  }

  public void checkPriority(CSQueue queue, int expectedPriority) {
    Assert.assertEquals(expectedPriority, queue.getPriority().getPriority());
  }

  public void checkReservedResource(CSQueue queue, String partition, int reserved) {
    ResourceUsage ru = queue.getQueueResourceUsage();
    Assert.assertEquals(reserved, ru.getReserved(partition).getMemorySize());
  }

  static class IsPreemptionRequestForQueueAndNode
      extends ArgumentMatcher<ContainerPreemptEvent> {
    private final ApplicationAttemptId appAttId;
    private final String queueName;
    private final NodeId nodeId;

    IsPreemptionRequestForQueueAndNode(ApplicationAttemptId appAttId,
        String queueName, NodeId nodeId) {
      this.appAttId = appAttId;
      this.queueName = queueName;
      this.nodeId = nodeId;
    }
    @Override
    public boolean matches(Object o) {
      ContainerPreemptEvent cpe = (ContainerPreemptEvent)o;

      return appAttId.equals(cpe.getAppId())
          && queueName.equals(cpe.getContainer().getQueueName())
          && nodeId.equals(cpe.getContainer().getAllocatedNode());
    }
    @Override
    public String toString() {
      return appAttId.toString();
    }
  }
}