/**
 * 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
 * <p>
 * http://www.apache.org/licenses/LICENSE-2.0
 * <p>
 * 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.scheduler.capacity;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
import org.apache.hadoop.yarn.api.records.ApplicationId;
import org.apache.hadoop.yarn.api.records.ContainerId;
import org.apache.hadoop.yarn.api.records.Priority;
import org.apache.hadoop.yarn.api.records.QueueState;
import org.apache.hadoop.yarn.api.records.Resource;
import org.apache.hadoop.yarn.api.records.ResourceRequest;
import org.apache.hadoop.yarn.conf.YarnConfiguration;
import org.apache.hadoop.yarn.factories.RecordFactory;
import org.apache.hadoop.yarn.factory.providers.RecordFactoryProvider;
import org.apache.hadoop.yarn.server.resourcemanager.MockNodes;
import org.apache.hadoop.yarn.server.resourcemanager.MockRM;
import org.apache.hadoop.yarn.server.resourcemanager.RMContext;
import org.apache.hadoop.yarn.server.resourcemanager.RMContextImpl;
import org.apache.hadoop.yarn.server.resourcemanager.placement
    .ApplicationPlacementContext;
import org.apache.hadoop.yarn.server.resourcemanager.rmapp.RMApp;
import org.apache.hadoop.yarn.server.resourcemanager.rmapp.RMAppEvent;
import org.apache.hadoop.yarn.server.resourcemanager.rmapp.RMAppEventType;
import org.apache.hadoop.yarn.server.resourcemanager.rmapp.RMAppState;
import org.apache.hadoop.yarn.server.resourcemanager.rmapp.attempt
    .RMAppAttemptState;
import org.apache.hadoop.yarn.server.resourcemanager.rmnode.RMNode;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler
    .SchedulerDynamicEditException;

import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity
    .queuemanagement.GuaranteedOrZeroCapacityOverTimePolicy;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.common
    .QueueEntitlement;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.event
    .AppAddedSchedulerEvent;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.event
    .NodeAddedSchedulerEvent;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.event
    .SchedulerEvent;

import org.apache.hadoop.yarn.server.resourcemanager.scheduler.policy
    .FairOrderingPolicy;
import org.apache.hadoop.yarn.server.resourcemanager.security
    .ClientToAMTokenSecretManagerInRM;
import org.apache.hadoop.yarn.server.resourcemanager.security
    .NMTokenSecretManagerInRM;
import org.apache.hadoop.yarn.server.resourcemanager.security
    .RMContainerTokenSecretManager;
import org.apache.hadoop.yarn.server.utils.BuilderUtils;
import org.apache.hadoop.yarn.util.resource.Resources;
import org.junit.Assert;
import org.junit.Test;

import java.io.IOException;
import java.util.Collections;
import java.util.List;

import static org.apache.hadoop.yarn.server.resourcemanager.placement
    .UserGroupMappingPlacementRule.CURRENT_USER_MAPPING;
import static org.apache.hadoop.yarn.server.resourcemanager.scheduler
    .capacity.CSQueueUtils.EPSILON;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

/**
 * Tests for creation and reinitialization of auto created leaf queues
 * under a ManagedParentQueue.
 */
public class TestCapacitySchedulerAutoQueueCreation
    extends TestCapacitySchedulerAutoCreatedQueueBase {

  private static final Log LOG = LogFactory.getLog(
      TestCapacitySchedulerAutoQueueCreation.class);

  private static final Resource TEMPLATE_MAX_RES = Resource.newInstance(16 *
          GB,
      48);
  private static final Resource TEMPLATE_MIN_RES = Resource.newInstance(1638,
      4);


  @Test(timeout = 10000)
  public void testAutoCreateLeafQueueCreation() throws Exception {

    try {
      // submit an app
      submitApp(mockRM, cs.getQueue(PARENT_QUEUE), USER0, USER0, 1, 1);

      // check preconditions
      List<ApplicationAttemptId> appsInC = cs.getAppsInQueue(PARENT_QUEUE);
      assertEquals(1, appsInC.size());
      assertNotNull(cs.getQueue(USER0));

      AutoCreatedLeafQueue autoCreatedLeafQueue =
          (AutoCreatedLeafQueue) cs.getQueue(USER0);
      ManagedParentQueue parentQueue = (ManagedParentQueue) cs.getQueue(
          PARENT_QUEUE);
      assertEquals(parentQueue, autoCreatedLeafQueue.getParent());
      validateInitialQueueEntitlement(parentQueue, USER0, 0.1f);
      validateUserAndAppLimits(autoCreatedLeafQueue, 1000, 1000);

      assertTrue(autoCreatedLeafQueue
          .getOrderingPolicy() instanceof FairOrderingPolicy);
    } finally {
      cleanupQueue(USER0);
    }
  }

  @Test
  public void testReinitializeStoppedAutoCreatedLeafQueue() throws Exception {
    try {
      String host = "127.0.0.1";
      RMNode node = MockNodes.newNodeInfo(0, MockNodes.newResource(4 * GB), 1,
          host);
      cs.handle(new NodeAddedSchedulerEvent(node));

      // submit an app

      RMApp app1 = mockRM.submitApp(GB, "test-auto-queue-creation-1", USER0,
          null, USER0);

      RMApp app2 = mockRM.submitApp(GB, "test-auto-queue-creation-2", USER1,
          null, USER1);
      // check preconditions
      List<ApplicationAttemptId> appsInC = cs.getAppsInQueue(PARENT_QUEUE);
      assertEquals(2, appsInC.size());

      assertNotNull(cs.getQueue(USER0));
      assertNotNull(cs.getQueue(USER1));

      AutoCreatedLeafQueue user0Queue = (AutoCreatedLeafQueue) cs.getQueue(
          USER0);
      AutoCreatedLeafQueue user1Queue = (AutoCreatedLeafQueue) cs.getQueue(
          USER0);
      ManagedParentQueue parentQueue = (ManagedParentQueue) cs.getQueue(
          PARENT_QUEUE);
      assertEquals(parentQueue, user0Queue.getParent());
      assertEquals(parentQueue, user1Queue.getParent());
      validateInitialQueueEntitlement(parentQueue, USER0, 0.2f);
      validateInitialQueueEntitlement(parentQueue, USER1, 0.2f);

      ApplicationAttemptId appAttemptId = appsInC.get(0);

      Priority priority = TestUtils.createMockPriority(1);
      RecordFactory recordFactory = RecordFactoryProvider.getRecordFactory(
          null);
      ResourceRequest r1 = TestUtils.createResourceRequest(ResourceRequest.ANY,
          1 * GB, 1, true, priority, recordFactory);

      cs.allocate(appAttemptId, Collections.<ResourceRequest>singletonList(r1),
          Collections.<ContainerId>emptyList(), Collections.singletonList(host),
          null, NULL_UPDATE_REQUESTS);

      //And this will result in container assignment for app1
      CapacityScheduler.schedule(cs);

      //change state to draining
      user0Queue.stopQueue();

      cs.killAllAppsInQueue(USER0);

      mockRM.waitForState(appAttemptId, RMAppAttemptState.KILLED);

      mockRM.waitForState(appAttemptId.getApplicationId(), RMAppState.KILLED);

      //change state to stopped
      user0Queue.stopQueue();
      assertEquals(QueueState.STOPPED,
          user0Queue.getQueueInfo().getQueueState());

      cs.reinitialize(cs.getConf(), mockRM.getRMContext());

      AutoCreatedLeafQueue user0QueueReinited =
          (AutoCreatedLeafQueue) cs.getQueue(USER0);

      validateCapacities(user0QueueReinited, 0.0f, 0.0f, 1.0f, 1.0f);

      AutoCreatedLeafQueue leafQueue = (AutoCreatedLeafQueue) cs.getQueue(
          USER1);
      validateInitialQueueEntitlement(parentQueue, leafQueue.getQueueName(),
          0.1f);

    } finally {
      cleanupQueue(USER0);
    }
  }

  @Test
  public void testConvertAutoCreateDisabledOnManagedParentQueueFails()
      throws Exception {
    CapacityScheduler newCS = new CapacityScheduler();
    try {
      CapacitySchedulerConfiguration newConf =
          new CapacitySchedulerConfiguration();
      setupQueueConfiguration(newConf);

      newConf.setAutoCreateChildQueueEnabled(C, false);

      newCS.setConf(new YarnConfiguration());
      newCS.setRMContext(mockRM.getRMContext());
      newCS.init(cs.getConf());
      newCS.start();

      newCS.reinitialize(newConf,
          new RMContextImpl(null, null, null, null, null, null,
              new RMContainerTokenSecretManager(newConf),
              new NMTokenSecretManagerInRM(newConf),
              new ClientToAMTokenSecretManagerInRM(), null));

    } catch (IOException e) {
      //expected exception
    } finally {
      newCS.stop();
    }
  }

  @Test
  public void testConvertLeafQueueToParentQueueWithAutoCreate()
      throws Exception {
    CapacityScheduler newCS = new CapacityScheduler();
    try {
      CapacitySchedulerConfiguration newConf =
          new CapacitySchedulerConfiguration();
      setupQueueConfiguration(newConf);
      newConf.setAutoCreatedLeafQueueConfigCapacity(A1, A1_CAPACITY / 10);
      newConf.setAutoCreateChildQueueEnabled(A1, true);

      newCS.setConf(new YarnConfiguration());
      newCS.setRMContext(mockRM.getRMContext());
      newCS.init(cs.getConf());
      newCS.start();

      final LeafQueue a1Queue = (LeafQueue) newCS.getQueue("a1");
      a1Queue.stopQueue();

      newCS.reinitialize(newConf,
          new RMContextImpl(null, null, null, null, null, null,
              new RMContainerTokenSecretManager(newConf),
              new NMTokenSecretManagerInRM(newConf),
              new ClientToAMTokenSecretManagerInRM(), null));

    } finally {
      newCS.stop();
    }
  }

  @Test
  public void testConvertFailsFromParentQueueToManagedParentQueue()
      throws Exception {
    CapacityScheduler newCS = new CapacityScheduler();
    try {
      CapacitySchedulerConfiguration newConf =
          new CapacitySchedulerConfiguration();
      setupQueueConfiguration(newConf);
      newConf.setAutoCreatedLeafQueueConfigCapacity(A, A_CAPACITY / 10);
      newConf.setAutoCreateChildQueueEnabled(A, true);

      newCS.setConf(new YarnConfiguration());
      newCS.setRMContext(mockRM.getRMContext());
      newCS.init(cs.getConf());
      newCS.start();

      final ParentQueue a1Queue = (ParentQueue) newCS.getQueue("a");
      a1Queue.stopQueue();

      newCS.reinitialize(newConf,
          new RMContextImpl(null, null, null, null, null, null,
              new RMContainerTokenSecretManager(newConf),
              new NMTokenSecretManagerInRM(newConf),
              new ClientToAMTokenSecretManagerInRM(), null));

      fail("Expected exception while converting a parent queue to"
          + " an auto create enabled parent queue");
    } catch (IOException e) {
      //expected exception
    } finally {
      newCS.stop();
    }
  }

  @Test(timeout = 10000)
  public void testAutoCreateLeafQueueFailsWithNoQueueMapping()
      throws Exception {

    final String INVALID_USER = "invalid_user";

    // submit an app under a different queue name which does not exist
    // and queue mapping does not exist for this user
    RMApp app = mockRM.submitApp(GB, "app", INVALID_USER, null, INVALID_USER,
        false);
    mockRM.drainEvents();
    mockRM.waitForState(app.getApplicationId(), RMAppState.FAILED);
    assertEquals(RMAppState.FAILED, app.getState());
  }

  @Test(timeout = 10000)
  public void testQueueMappingValidationFailsWithInvalidParentQueueInMapping()
      throws Exception {

    MockRM newMockRM = setupSchedulerInstance();
    try {
      CapacityScheduler newCS =
          (CapacityScheduler) newMockRM.getResourceScheduler();

      //"a" is not auto create enabled

      //dynamic queue mapping
      try {
        setupQueueMapping(newCS, CURRENT_USER_MAPPING, "a",
            CURRENT_USER_MAPPING);
        newCS.updatePlacementRules();
        fail("Expected invalid parent queue mapping failure");

      } catch (IOException e) {
        //expected exception
        assertTrue(e.getMessage().contains(
            "invalid parent queue which does not have auto creation of leaf "
                + "queues enabled [" + "a" + "]"));
      }

      //"a" is not auto create enabled and app_user does not exist as a leaf
      // queue
      //static queue mapping
      try {
        setupQueueMapping(newCS, "app_user", "INVALID_PARENT_QUEUE",
            "app_user");
        newCS.updatePlacementRules();
        fail("Expected invalid parent queue mapping failure");
      } catch (IOException e) {
        //expected exception
        assertTrue(e.getMessage()
            .contains("invalid parent queue [" + "INVALID_PARENT_QUEUE" + "]"));
      }
    } finally {
      if (newMockRM != null) {
        ((CapacityScheduler) newMockRM.getResourceScheduler()).stop();
        newMockRM.stop();
      }
    }
  }

  @Test(timeout = 10000)
  public void testQueueMappingUpdatesFailsOnRemovalOfParentQueueInMapping()
      throws Exception {

    MockRM newMockRM = setupSchedulerInstance();

    try {
      CapacityScheduler newCS =
          (CapacityScheduler) newMockRM.getResourceScheduler();

      setupQueueMapping(newCS, CURRENT_USER_MAPPING, "c", CURRENT_USER_MAPPING);
      newCS.updatePlacementRules();

      try {
        setupQueueMapping(newCS, CURRENT_USER_MAPPING, "",
            CURRENT_USER_MAPPING);
        newCS.updatePlacementRules();
        fail("Expected invalid parent queue mapping failure");
      } catch (IOException e) {
        //expected exception
        assertTrue(e.getMessage().contains("invalid parent queue []"));
      }
    } finally {
      if (newMockRM != null) {
        ((CapacityScheduler) newMockRM.getResourceScheduler()).stop();
        newMockRM.stop();
      }
    }
  }

  @Test
  public void testParentQueueUpdateInQueueMappingFailsAfterAutoCreation()
      throws Exception {

    MockRM newMockRM = setupSchedulerInstance();
    CapacityScheduler newCS =
        (CapacityScheduler) newMockRM.getResourceScheduler();

    try {
      submitApp(newCS, USER0, USER0, PARENT_QUEUE);

      assertNotNull(newCS.getQueue(USER0));

      setupQueueMapping(newCS, USER0, "d", USER0);
      newCS.updatePlacementRules();

      RMContext rmContext = mock(RMContext.class);
      when(rmContext.getDispatcher()).thenReturn(dispatcher);
      newCS.setRMContext(rmContext);

      ApplicationId appId = BuilderUtils.newApplicationId(1, 1);
      SchedulerEvent addAppEvent = new AppAddedSchedulerEvent(appId, USER0,
          USER0, new ApplicationPlacementContext(USER0, "d"));
      newCS.handle(addAppEvent);

      RMAppEvent event = new RMAppEvent(appId, RMAppEventType.APP_REJECTED,
          "error");
      dispatcher.spyOnNextEvent(event, 10000);
    } finally {
      if (newMockRM != null) {
        ((CapacityScheduler) newMockRM.getResourceScheduler()).stop();
        newMockRM.stop();
      }
    }
  }

  @Test
  public void testAutoCreationFailsWhenParentCapacityExceeded()
      throws Exception {
    MockRM newMockRM = setupSchedulerInstance();
    CapacityScheduler newCS =
        (CapacityScheduler) newMockRM.getResourceScheduler();

    try {
      CapacitySchedulerConfiguration conf = newCS.getConfiguration();
      conf.setShouldFailAutoQueueCreationWhenGuaranteedCapacityExceeded(C,
          true);

      newCS.reinitialize(conf, newMockRM.getRMContext());

      // Test add one auto created queue dynamically and manually modify
      // capacity
      ManagedParentQueue parentQueue = (ManagedParentQueue) newCS.getQueue("c");
      AutoCreatedLeafQueue c1 = new AutoCreatedLeafQueue(newCS, "c1",
          parentQueue);
      newCS.addQueue(c1);
      c1.setCapacity(0.5f);
      c1.setAbsoluteCapacity(c1.getParent().getAbsoluteCapacity() * 1f);
      c1.setMaxCapacity(1f);

      setEntitlement(c1, new QueueEntitlement(0.5f, 1f));

      AutoCreatedLeafQueue c2 = new AutoCreatedLeafQueue(newCS, "c2",
          parentQueue);
      newCS.addQueue(c2);
      setEntitlement(c2, new QueueEntitlement(0.5f, 1f));

      try {
        AutoCreatedLeafQueue c3 = new AutoCreatedLeafQueue(newCS, "c3",
            parentQueue);
        newCS.addQueue(c3);
        fail("Expected exception for auto queue creation failure");
      } catch (SchedulerDynamicEditException e) {
        //expected exception
      }
    } finally {
      if (newMockRM != null) {
        ((CapacityScheduler) newMockRM.getResourceScheduler()).stop();
        newMockRM.stop();
      }
    }
  }

  @Test
  public void testAutoCreatedQueueActivationDeactivation() throws Exception {

    try {
      CSQueue parentQueue = cs.getQueue(PARENT_QUEUE);

      //submit app1 as USER1
      submitApp(mockRM, parentQueue, USER1, USER1, 1, 1);
      validateInitialQueueEntitlement(parentQueue, USER1, 0.1f);

      //submit another app2 as USER2
      ApplicationId user2AppId = submitApp(mockRM, parentQueue, USER2, USER2, 2,
          1);
      validateInitialQueueEntitlement(parentQueue, USER2, 0.2f);

      //submit another app3 as USER1
      submitApp(mockRM, parentQueue, USER1, USER1, 3, 2);

      //validate total activated abs capacity remains the same
      GuaranteedOrZeroCapacityOverTimePolicy autoCreatedQueueManagementPolicy =
          (GuaranteedOrZeroCapacityOverTimePolicy) ((ManagedParentQueue)
              parentQueue)
              .getAutoCreatedQueueManagementPolicy();
      assertEquals(autoCreatedQueueManagementPolicy
          .getAbsoluteActivatedChildQueueCapacity(), 0.2f, EPSILON);

      //submit user_3 app. This cant be scheduled since there is no capacity
      submitApp(mockRM, parentQueue, USER3, USER3, 4, 1);
      final CSQueue user3LeafQueue = cs.getQueue(USER3);
      validateCapacities((AutoCreatedLeafQueue) user3LeafQueue, 0.0f, 0.0f,
          1.0f, 1.0f);

      assertEquals(autoCreatedQueueManagementPolicy
          .getAbsoluteActivatedChildQueueCapacity(), 0.2f, EPSILON);

      //deactivate USER2 queue
      cs.killAllAppsInQueue(USER2);
      mockRM.waitForState(user2AppId, RMAppState.KILLED);

      //Verify if USER_2 can be deactivated since it has no pending appsA
      List<QueueManagementChange> queueManagementChanges =
          autoCreatedQueueManagementPolicy.computeQueueManagementChanges();

      ManagedParentQueue managedParentQueue = (ManagedParentQueue) parentQueue;
      managedParentQueue.validateAndApplyQueueManagementChanges(
          queueManagementChanges);

      validateDeactivatedQueueEntitlement(parentQueue, USER2, 0.2f,
          queueManagementChanges);

      //USER_3 should now get activated
      validateActivatedQueueEntitlement(parentQueue, USER3, 0.2f,
          queueManagementChanges);

    } finally {
      cleanupQueue(USER1);
      cleanupQueue(USER2);
      cleanupQueue(USER3);
    }
  }

  @Test
  public void testClusterResourceUpdationOnAutoCreatedLeafQueues() throws
      Exception {

    MockRM newMockRM = setupSchedulerInstance();
    try {
      CapacityScheduler newCS =
          (CapacityScheduler) newMockRM.getResourceScheduler();

      CSQueue parentQueue = newCS.getQueue(PARENT_QUEUE);

      //submit app1 as USER1
      submitApp(newMockRM, parentQueue, USER1, USER1, 1, 1);
      validateInitialQueueEntitlement(newCS, parentQueue, USER1, 0.1f);
      CSQueue user1LeafQueue = newCS.getQueue(USER1);

      //submit another app2 as USER2
      submitApp(newMockRM, parentQueue, USER2, USER2, 2, 1);
      validateInitialQueueEntitlement(newCS, parentQueue, USER2, 0.2f);
      CSQueue user2LeafQueue = newCS.getQueue(USER2);

      //validate total activated abs capacity remains the same
      GuaranteedOrZeroCapacityOverTimePolicy autoCreatedQueueManagementPolicy =
          (GuaranteedOrZeroCapacityOverTimePolicy) ((ManagedParentQueue)
              parentQueue)
              .getAutoCreatedQueueManagementPolicy();
      assertEquals(autoCreatedQueueManagementPolicy
          .getAbsoluteActivatedChildQueueCapacity(), 0.2f, EPSILON);

      //submit user_3 app. This cant be scheduled since there is no capacity
      submitApp(newMockRM, parentQueue, USER3, USER3, 3, 1);
      final CSQueue user3LeafQueue = newCS.getQueue(USER3);
      validateCapacities((AutoCreatedLeafQueue) user3LeafQueue, 0.0f, 0.0f,
          1.0f, 1.0f);

      assertEquals(autoCreatedQueueManagementPolicy
          .getAbsoluteActivatedChildQueueCapacity(), 0.2f, EPSILON);

      // add new NM.
      newMockRM.registerNode("127.0.0.3:1234", 125 * GB, 20);

      // There will be change in effective resource when nodes are added
      // since we deal with percentages

      Resource MAX_RES = Resources.addTo(TEMPLATE_MAX_RES,
          Resources.createResource(125 * GB, 20));

      Resource MIN_RES = Resources.createResource(14438, 6);

      Assert.assertEquals("Effective Min resource for USER3 is not correct",
          Resources.none(),
          user3LeafQueue.getQueueResourceQuotas().getEffectiveMinResource());
      Assert.assertEquals("Effective Max resource for USER3 is not correct",
          MAX_RES,
          user3LeafQueue.getQueueResourceQuotas().getEffectiveMaxResource());

      Assert.assertEquals("Effective Min resource for USER2 is not correct",
          MIN_RES,
          user1LeafQueue.getQueueResourceQuotas().getEffectiveMinResource());
      Assert.assertEquals("Effective Max resource for USER2 is not correct",
          MAX_RES,
          user1LeafQueue.getQueueResourceQuotas().getEffectiveMaxResource());

      Assert.assertEquals("Effective Min resource for USER1 is not correct",
          MIN_RES,
          user2LeafQueue.getQueueResourceQuotas().getEffectiveMinResource());
      Assert.assertEquals("Effective Max resource for USER1 is not correct",
          MAX_RES,
          user2LeafQueue.getQueueResourceQuotas().getEffectiveMaxResource());

      // unregister one NM.
      newMockRM.unRegisterNode(nm3);
      Resource MIN_RES_UPDATED = Resources.createResource(12800, 2);
      Resource MAX_RES_UPDATED = Resources.createResource(128000, 20);

      // After loosing one NM, resources will reduce
      Assert.assertEquals("Effective Min resource for USER2 is not correct",
          MIN_RES_UPDATED,
          user1LeafQueue.getQueueResourceQuotas().getEffectiveMinResource());
      Assert.assertEquals("Effective Max resource for USER2 is not correct",
          MAX_RES_UPDATED,
          user2LeafQueue.getQueueResourceQuotas().getEffectiveMaxResource());

    } finally {
      cleanupQueue(USER1);
      cleanupQueue(USER2);
      cleanupQueue(USER3);
      if (newMockRM != null) {
        ((CapacityScheduler) newMockRM.getResourceScheduler()).stop();
        newMockRM.stop();
      }
    }
  }

  @Test
  public void testAutoCreatedQueueInheritsNodeLabels() throws Exception {

    try {
      String host = "127.0.0.1";
      RMNode node = MockNodes.newNodeInfo(0, MockNodes.newResource(4 * GB), 1,
          host);
      cs.handle(new NodeAddedSchedulerEvent(node));

      CSQueue parentQueue = cs.getQueue(PARENT_QUEUE);

      submitApp(USER1, USER1, NODEL_LABEL_GPU);
      //submit app1 as USER1
      validateInitialQueueEntitlement(parentQueue, USER1, 0.1f);
    } finally {
      cleanupQueue(USER1);
    }
  }

  @Test
  public void testReinitializeQueuesWithAutoCreatedLeafQueues()
      throws Exception {

    MockRM newMockRM = setupSchedulerInstance();
    try {
      CapacityScheduler newCS =
          (CapacityScheduler) newMockRM.getResourceScheduler();
      CapacitySchedulerConfiguration conf = newCS.getConfiguration();

      CSQueue parentQueue = newCS.getQueue(PARENT_QUEUE);

      //submit app1 as USER1
      submitApp(newMockRM, parentQueue, USER1, USER1, 1, 1);
      validateInitialQueueEntitlement(newCS, parentQueue, USER1, 0.1f);

      //submit another app2 as USER2
      ApplicationId user2AppId = submitApp(newMockRM, parentQueue, USER2, USER2,
          2, 1);
      validateInitialQueueEntitlement(newCS, parentQueue, USER2, 0.2f);

      //update parent queue capacity
      conf.setCapacity(C, 30f);
      conf.setCapacity(D, 10f);
      conf.setMaximumCapacity(C, 50f);

      newCS.reinitialize(conf, newMockRM.getRMContext());

      // validate that leaf queues abs capacity is now changed
      AutoCreatedLeafQueue user0Queue = (AutoCreatedLeafQueue) newCS.getQueue(
          USER1);
      validateCapacities(user0Queue, 0.5f, 0.15f, 1.0f, 0.5f);
      validateUserAndAppLimits(user0Queue, 1500, 1500);

      //update leaf queue template capacities
      conf.setAutoCreatedLeafQueueConfigCapacity(C, 30f);
      conf.setAutoCreatedLeafQueueConfigMaxCapacity(C, 40f);

      newCS.reinitialize(conf, newMockRM.getRMContext());
      validateCapacities(user0Queue, 0.3f, 0.09f, 0.4f, 0.2f);
      validateUserAndAppLimits(user0Queue, 900, 900);

      //submit app1 as USER3
      submitApp(newMockRM, parentQueue, USER3, USER3, 3, 1);
      validateInitialQueueEntitlement(newCS, parentQueue, USER3, 0.27f);
      AutoCreatedLeafQueue user3Queue = (AutoCreatedLeafQueue) newCS.getQueue(
          USER1);
      validateCapacities(user3Queue, 0.3f, 0.09f, 0.4f, 0.2f);
      validateUserAndAppLimits(user3Queue, 900, 900);

      //submit app1 as USER1 - is already activated. there should be no diff
      // in capacities
      submitApp(newMockRM, parentQueue, USER3, USER3, 4, 2);
      validateInitialQueueEntitlement(newCS, parentQueue, USER3, 0.27f);
      validateCapacities(user3Queue, 0.3f, 0.09f, 0.4f, 0.2f);
      validateUserAndAppLimits(user3Queue, 900, 900);

    } finally {
      cleanupQueue(USER1);
      cleanupQueue(USER2);
      if (newMockRM != null) {
        ((CapacityScheduler) newMockRM.getResourceScheduler()).stop();
        newMockRM.stop();
      }
    }
  }
}
