@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);
}