@Override
  public N getSchedulerNode(NodeId nodeId) {
    return nodeTracker.getNode(nodeId);
  }
@Test(timeout = 30000l)
  public void testContainerReleaseWithAllocationTags() throws Exception {
    // Currently only can be tested against capacity scheduler.
    if (getSchedulerType().equals(SchedulerType.CAPACITY)) {
      final String testTag1 = "some-tag";
      final String testTag2 = "some-other-tag";
      YarnConfiguration conf = getConf();
      conf.set(YarnConfiguration.RM_PLACEMENT_CONSTRAINTS_HANDLER, "scheduler");
      MockRM rm1 = new MockRM(conf);
      rm1.start();
      MockNM nm1 = new MockNM("127.0.0.1:1234",
          10240, rm1.getResourceTrackerService());
      nm1.registerNode();
      MockRMAppSubmissionData data =
          MockRMAppSubmissionData.Builder.createWithMemory(200, rm1)
          .withAppName("name")
          .withUser("user")
          .withAcls(new HashMap<>())
          .withUnmanagedAM(false)
          .withQueue("default")
          .withMaxAppAttempts(-1)
          .withCredentials(null)
          .withAppType("Test")
          .withWaitForAppAcceptedState(false)
          .withKeepContainers(true)
          .build();
      RMApp app1 =
          MockRMAppSubmitter.submit(rm1, data);
      MockAM am1 = MockRM.launchAndRegisterAM(app1, rm1, nm1);

      // allocate 1 container with tag1
      SchedulingRequest sr = SchedulingRequest
          .newInstance(1l, Priority.newInstance(1),
              ExecutionTypeRequest.newInstance(ExecutionType.GUARANTEED),
              Sets.newHashSet(testTag1),
              ResourceSizing.newInstance(1, Resource.newInstance(1024, 1)),
              null);

      // allocate 3 containers with tag2
      SchedulingRequest sr1 = SchedulingRequest
          .newInstance(2l, Priority.newInstance(1),
              ExecutionTypeRequest.newInstance(ExecutionType.GUARANTEED),
              Sets.newHashSet(testTag2),
              ResourceSizing.newInstance(3, Resource.newInstance(1024, 1)),
              null);

      AllocateRequest ar = AllocateRequest.newBuilder()
          .schedulingRequests(Lists.newArrayList(sr, sr1)).build();
      am1.allocate(ar);
      nm1.nodeHeartbeat(true);

      List<Container> allocated = new ArrayList<>();
      while (allocated.size() < 4) {
        AllocateResponse rsp = am1
            .allocate(new ArrayList<>(), new ArrayList<>());
        allocated.addAll(rsp.getAllocatedContainers());
        nm1.nodeHeartbeat(true);
        Thread.sleep(1000);
      }

      Assert.assertEquals(4, allocated.size());

      Set<Container> containers = allocated.stream()
          .filter(container -> container.getAllocationRequestId() == 1l)
          .collect(Collectors.toSet());
      Assert.assertNotNull(containers);
      Assert.assertEquals(1, containers.size());
      ContainerId cid = containers.iterator().next().getId();

      // mock container start
      rm1.getRMContext().getScheduler()
          .getSchedulerNode(nm1.getNodeId()).containerStarted(cid);

      // verifies the allocation is made with correct number of tags
      Map<String, Long> nodeTags = rm1.getRMContext()
          .getAllocationTagsManager()
          .getAllocationTagsWithCount(nm1.getNodeId());
      Assert.assertNotNull(nodeTags.get(testTag1));
      Assert.assertEquals(1, nodeTags.get(testTag1).intValue());

      // release a container
      am1.allocate(new ArrayList<>(), Lists.newArrayList(cid));

      // before NM confirms, the tag should still exist
      nodeTags = rm1.getRMContext().getAllocationTagsManager()
          .getAllocationTagsWithCount(nm1.getNodeId());
      Assert.assertNotNull(nodeTags);
      Assert.assertNotNull(nodeTags.get(testTag1));
      Assert.assertEquals(1, nodeTags.get(testTag1).intValue());

      // NM reports back that container is released
      // RM should cleanup the tag
      ContainerStatus cs = ContainerStatus.newInstance(cid,
          ContainerState.COMPLETE, "", 0);
      nm1.nodeHeartbeat(Lists.newArrayList(cs), true);

      // Wait on condition
      // 1) tag1 doesn't exist anymore
      // 2) num of tag2 is still 3
      GenericTestUtils.waitFor(() -> {
        Map<String, Long> tags = rm1.getRMContext()
            .getAllocationTagsManager()
            .getAllocationTagsWithCount(nm1.getNodeId());
        return tags.get(testTag1) == null &&
            tags.get(testTag2).intValue() == 3;
      }, 500, 3000);
    }
  }