@Test
  public void testDistributedLeaderStart() throws Exception {
    RMContextImpl rmContext = new RMContextImpl();
    rmContext.setDistributedEnabled(true);
    rmContext.setHAEnabled(true);
    GroupMembershipService groupMembershipService = new GroupMembershipService(null, rmContext);
    groupMembershipService.init(conf);
    NMTokenSecretManagerInRM nmTokenSecretManager = new NMTokenSecretManagerInRM(conf, rmContext);
    try {
      groupMembershipService.start();

      rmContext.setRMGroupMembershipService(groupMembershipService);
      rmContext.setStateStore(new NDBRMStateStore());

      assertTrue("should be leader", rmContext.isLeader());
      nmTokenSecretManager.start();
      assertNotNull("did not roll master key", nmTokenSecretManager.getCurrentKey());

      RMStateStore.RMState state = rmContext.getStateStore().loadState(rmContext);
      assertNotNull(
          "key not persisted to the database",
          state.getSecretTokenMamagerKey(RMStateStore.KeyType.CURRENTNMTOKENMASTERKEY));

      NMTokenSecretManagerInRM nmTokenSecretManager2 =
          new NMTokenSecretManagerInRM(conf, rmContext);
      nmTokenSecretManager2.recover(state);

      assertEquals(nmTokenSecretManager.getCurrentKey(), nmTokenSecretManager2.getCurrentKey());
      assertEquals(nmTokenSecretManager.getNextKey(), nmTokenSecretManager2.getNextKey());
    } finally {
      groupMembershipService.stop();
      nmTokenSecretManager.stop();
      DefaultMetricsSystem.shutdown();
    }
  }
  @Test
  public void testNonLeaderKeyReception()
      throws InterruptedException, StorageInitializtionException, Exception {
    // create a groupMembershipService that will be leader
    RMContextImpl rmContext = new RMContextImpl();
    rmContext.setDistributedEnabled(true);
    rmContext.setHAEnabled(true);
    GroupMembershipService groupMembershipService = new GroupMembershipService(null, rmContext);
    groupMembershipService.init(conf);
    NMTokenSecretManagerInRM nmTokenSecretManager = new NMTokenSecretManagerInRM(conf, rmContext);
    // create a resrouce tracker
    conf.setBoolean(YarnConfiguration.RM_HA_ENABLED, true);
    conf.setBoolean(YarnConfiguration.DISTRIBUTED_RM, true);
    MockRM mockRM = new MockRM(conf);
    mockRM.init(conf);

    try {
      groupMembershipService.start();
      rmContext.setRMGroupMembershipService(groupMembershipService);
      rmContext.setStateStore(new NDBRMStateStore());

      while (!rmContext.isLeader()) {
        Thread.sleep(1000);
      }

      mockRM.start();

      if (mockRM.getRMContext().isDistributedEnabled() && !mockRM.getRMContext().isLeader()) {
        conf.set(
            YarnConfiguration.EVENT_RT_CONFIG_PATH, "target/test-classes/RT_EventAPIConfig.ini");
        NdbRtStreamingProcessor rtStreamingProcessor =
            new NdbRtStreamingProcessor(mockRM.getRMContext());
        RMStorageFactory.kickTheNdbEventStreamingAPI(false, conf);
        new Thread(rtStreamingProcessor).start();
      }

      // this should be a resource tracker not a scheduler
      assertFalse(mockRM.getRMContext().isLeader());

      // simulate creation of a token on the sheduler
      nmTokenSecretManager.start();

      assertNotNull("did not roll master key", nmTokenSecretManager.getCurrentKey());
      Thread.sleep(1000);
      dummyUpdate();
      Thread.sleep(1000);
      RMStateStore.RMState state = rmContext.getStateStore().loadState(rmContext);
      assertEquals(
          "key not persisted to the database",
          state.getSecretTokenMamagerKey(RMStateStore.KeyType.CURRENTNMTOKENMASTERKEY),
          nmTokenSecretManager.getCurrentKey());

      assertEquals(
          nmTokenSecretManager.getCurrentKey(),
          mockRM.getRMContext().getNMTokenSecretManager().getCurrentKey());
      assertEquals(
          nmTokenSecretManager.getNextKey(),
          mockRM.getRMContext().getNMTokenSecretManager().getNextKey());
    } finally {
      groupMembershipService.stop();
      mockRM.stop();
      nmTokenSecretManager.stop();
      DefaultMetricsSystem.shutdown();
    }
  }