@Test
public void testAMRMTokenUpdate() throws Exception {

  Configuration conf = new Configuration();
  ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(
      ApplicationId.newInstance(1, 1), 1);
  AMRMTokenIdentifier oldTokenId = new AMRMTokenIdentifier(attemptId, 1);
  AMRMTokenIdentifier newTokenId = new AMRMTokenIdentifier(attemptId, 2);
  Token<AMRMTokenIdentifier> oldToken = new Token<AMRMTokenIdentifier>(
      oldTokenId.getBytes(), "oldpassword".getBytes(), oldTokenId.getKind(),
      new Text());
  Token<AMRMTokenIdentifier> newToken = new Token<AMRMTokenIdentifier>(
      newTokenId.getBytes(), "newpassword".getBytes(), newTokenId.getKind(),
      new Text());
  MockScheduler scheduler = new MockScheduler();
  scheduler.amToken = newToken;
  final LocalContainerAllocator lca =
      new StubbedLocalContainerAllocator(scheduler);
  lca.init(conf);
  lca.start();
  UserGroupInformation testUgi = UserGroupInformation.createUserForTesting(
      "someuser", new String[0]);
  testUgi.addToken(oldToken);
  testUgi.doAs(new PrivilegedExceptionAction<Void>() {
        @Override
        public Void run() throws Exception {
          lca.heartbeat();
          return null;
        }
  });
  lca.close();
  // verify there is only one AMRM token in the UGI and it matches the
  // updated token from the RM
  int tokenCount = 0;
  Token<? extends TokenIdentifier> ugiToken = null;
  for (Token<? extends TokenIdentifier> token : testUgi.getTokens()) {
    if (AMRMTokenIdentifier.KIND_NAME.equals(token.getKind())) {
      ugiToken = token;
      ++tokenCount;
    }
  }
  Assert.assertEquals("too many AMRM tokens", 1, tokenCount);
  Assert.assertArrayEquals("token identifier not updated",
      newToken.getIdentifier(), ugiToken.getIdentifier());
  Assert.assertArrayEquals("token password not updated",
      newToken.getPassword(), ugiToken.getPassword());
  Assert.assertEquals("AMRM token service not updated",
      new Text(ClientRMProxy.getAMRMTokenService(conf)),
      ugiToken.getService());
}