@Test(timeout = 15000)
  public void testRemoveExpiredMasterKeyInRMStateStore() throws Exception {
    MemoryRMStateStore memStore = new MemoryRMStateStore();
    memStore.init(conf);
    RMState rmState = memStore.getState();

    Set<DelegationKey> rmDTMasterKeyState = rmState.getRMDTSecretManagerState().getMasterKeyState();

    MockRM rm1 = new MyMockRM(conf, memStore);
    rm1.start();
    RMDelegationTokenSecretManager dtSecretManager = rm1.getRMDTSecretManager();

    // assert all master keys are saved
    Assert.assertEquals(dtSecretManager.getAllMasterKeys(), rmDTMasterKeyState);
    Set<DelegationKey> expiringKeys = new HashSet<DelegationKey>();
    expiringKeys.addAll(dtSecretManager.getAllMasterKeys());

    // wait for expiringKeys to expire
    while (true) {
      boolean allExpired = true;
      for (DelegationKey key : expiringKeys) {
        if (rmDTMasterKeyState.contains(key)) {
          allExpired = false;
        }
      }
      if (allExpired) break;
      Thread.sleep(500);
    }
  }
  @Test(timeout = 15000)
  public void testRMDTMasterKeyStateOnRollingMasterKey() throws Exception {
    MemoryRMStateStore memStore = new MemoryRMStateStore();
    memStore.init(conf);
    RMState rmState = memStore.getState();

    Map<RMDelegationTokenIdentifier, Long> rmDTState =
        rmState.getRMDTSecretManagerState().getTokenState();
    Set<DelegationKey> rmDTMasterKeyState = rmState.getRMDTSecretManagerState().getMasterKeyState();

    MockRM rm1 = new MyMockRM(conf, memStore);
    rm1.start();
    // on rm start, two master keys are created.
    // One is created at RMDTSecretMgr.startThreads.updateCurrentKey();
    // the other is created on the first run of
    // tokenRemoverThread.rollMasterKey()

    RMDelegationTokenSecretManager dtSecretManager = rm1.getRMDTSecretManager();
    // assert all master keys are saved
    Assert.assertEquals(dtSecretManager.getAllMasterKeys(), rmDTMasterKeyState);
    Set<DelegationKey> expiringKeys = new HashSet<DelegationKey>();
    expiringKeys.addAll(dtSecretManager.getAllMasterKeys());

    // record the current key
    DelegationKey oldCurrentKey =
        ((TestRMDelegationTokenSecretManager) dtSecretManager).getCurrentKey();

    // request to generate a RMDelegationToken
    GetDelegationTokenRequest request = mock(GetDelegationTokenRequest.class);
    when(request.getRenewer()).thenReturn("renewer1");
    GetDelegationTokenResponse response = rm1.getClientRMService().getDelegationToken(request);
    org.apache.hadoop.yarn.api.records.Token delegationToken = response.getRMDelegationToken();
    Token<RMDelegationTokenIdentifier> token1 =
        ConverterUtils.convertFromYarn(delegationToken, null);
    RMDelegationTokenIdentifier dtId1 = token1.decodeIdentifier();

    // wait for the first rollMasterKey
    while (((TestRMDelegationTokenSecretManager) dtSecretManager).numUpdatedKeys.get() < 1) {
      Thread.sleep(200);
    }

    // assert old-current-key and new-current-key exist
    Assert.assertTrue(rmDTMasterKeyState.contains(oldCurrentKey));
    DelegationKey newCurrentKey =
        ((TestRMDelegationTokenSecretManager) dtSecretManager).getCurrentKey();
    Assert.assertTrue(rmDTMasterKeyState.contains(newCurrentKey));

    // wait for token to expire
    // rollMasterKey is called every 1 second.
    while (((TestRMDelegationTokenSecretManager) dtSecretManager).numUpdatedKeys.get() < 6) {
      Thread.sleep(200);
    }

    Assert.assertFalse(rmDTState.containsKey(dtId1));
    rm1.stop();
  }
 @Override
 public void recover(RMState state) throws Exception {
   RMStateStore store = rmContext.getStateStore();
   assert store != null;
   // recover applications
   Map<ApplicationId, ApplicationState> appStates = state.getApplicationState();
   LOG.info("Recovering " + appStates.size() + " applications");
   for (ApplicationState appState : appStates.values()) {
     recoverApplication(appState, state);
   }
 }