public static List<RMNode> queryRMNodes(RMContext context,
      EnumSet<NodeState> acceptedStates) {
    // nodes contains nodes that are NEW, RUNNING, UNHEALTHY or DECOMMISSIONING.
    ArrayList<RMNode> results = new ArrayList<RMNode>();
    boolean hasActive = false;
    boolean hasInactive = false;
    for (NodeState nodeState : acceptedStates) {
      if (!hasInactive && nodeState.isInactiveState()) {
        hasInactive = true;
      }
      if (!hasActive && nodeState.isActiveState()) {
        hasActive = true;
      }
      if (hasActive && hasInactive) {
        break;
      }
    }
    if (hasActive) {
      for (RMNode rmNode : context.getRMNodes().values()) {
        if (acceptedStates.contains(rmNode.getState())) {
          results.add(rmNode);
        }
      }
    }

    // inactiveNodes contains nodes that are DECOMMISSIONED, LOST, OR REBOOTED
    if (hasInactive) {
      for (RMNode rmNode : context.getInactiveRMNodes().values()) {
        if ((rmNode != null) && acceptedStates.contains(rmNode.getState())) {
          results.add(rmNode);
        }
      }
    }
    return results;
  }
@Test
  public void testQueryRMNodes() throws Exception {
    RMContext rmContext = mock(RMContext.class);
    NodeId node1 = NodeId.newInstance("node1", 1234);
    RMNode rmNode1 = mock(RMNode.class);
    ConcurrentMap<NodeId, RMNode> inactiveList =
        new ConcurrentHashMap<NodeId, RMNode>();
    when(rmNode1.getState()).thenReturn(NodeState.SHUTDOWN);
    inactiveList.put(node1, rmNode1);
    when(rmContext.getInactiveRMNodes()).thenReturn(inactiveList);
    List<RMNode> result = RMServerUtils.queryRMNodes(rmContext,
        EnumSet.of(NodeState.SHUTDOWN));
    Assert.assertTrue(result.size() != 0);
    assertThat(result.get(0)).isEqualTo(rmNode1);
    when(rmNode1.getState()).thenReturn(NodeState.DECOMMISSIONED);
    result = RMServerUtils.queryRMNodes(rmContext,
        EnumSet.of(NodeState.DECOMMISSIONED));
    Assert.assertTrue(result.size() != 0);
    assertThat(result.get(0)).isEqualTo(rmNode1);
    when(rmNode1.getState()).thenReturn(NodeState.LOST);
    result = RMServerUtils.queryRMNodes(rmContext,
        EnumSet.of(NodeState.LOST));
    Assert.assertTrue(result.size() != 0);
    assertThat(result.get(0)).isEqualTo(rmNode1);
    when(rmNode1.getState()).thenReturn(NodeState.REBOOTED);
    result = RMServerUtils.queryRMNodes(rmContext,
        EnumSet.of(NodeState.REBOOTED));
    Assert.assertTrue(result.size() != 0);
    assertThat(result.get(0)).isEqualTo(rmNode1);
  }