/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.hadoop.yarn.server.resourcemanager.recovery;

import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.when;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;

import javax.crypto.SecretKey;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.ipc.CallerContext;
import org.apache.hadoop.security.token.Token;
import org.apache.hadoop.security.token.delegation.DelegationKey;
import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
import org.apache.hadoop.yarn.api.records.ApplicationId;
import org.apache.hadoop.yarn.api.records.ApplicationSubmissionContext;
import org.apache.hadoop.yarn.api.records.Container;
import org.apache.hadoop.yarn.api.records.ContainerId;
import org.apache.hadoop.yarn.api.records.FinalApplicationStatus;
import org.apache.hadoop.yarn.api.records.ReservationDefinition;
import org.apache.hadoop.yarn.api.records.ReservationId;
import org.apache.hadoop.yarn.api.records.Resource;
import org.apache.hadoop.yarn.api.records.impl.pb.ApplicationSubmissionContextPBImpl;
import org.apache.hadoop.yarn.api.records.impl.pb.ContainerLaunchContextPBImpl;
import org.apache.hadoop.yarn.api.records.impl.pb.ContainerPBImpl;
import org.apache.hadoop.yarn.conf.YarnConfiguration;
import org.apache.hadoop.yarn.event.Dispatcher;
import org.apache.hadoop.yarn.event.Event;
import org.apache.hadoop.yarn.event.EventHandler;
import org.apache.hadoop.yarn.proto.YarnProtos.ReservationAllocationStateProto;
import org.apache.hadoop.yarn.security.AMRMTokenIdentifier;
import org.apache.hadoop.yarn.security.client.RMDelegationTokenIdentifier;
import org.apache.hadoop.yarn.server.records.Version;
import org.apache.hadoop.yarn.server.resourcemanager.RMContext;
import org.apache.hadoop.yarn.server.resourcemanager.recovery.RMStateStore.RMDTSecretManagerState;
import org.apache.hadoop.yarn.server.resourcemanager.recovery.RMStateStore.RMState;
import org.apache.hadoop.yarn.server.resourcemanager.recovery.records.AMRMTokenSecretManagerState;
import org.apache.hadoop.yarn.server.resourcemanager.recovery.records.ApplicationAttemptStateData;
import org.apache.hadoop.yarn.server.resourcemanager.recovery.records.ApplicationStateData;
import org.apache.hadoop.yarn.server.resourcemanager.reservation.InMemoryReservationAllocation;
import org.apache.hadoop.yarn.server.resourcemanager.reservation.ReservationAllocation;
import org.apache.hadoop.yarn.server.resourcemanager.reservation.ReservationSystemTestUtil;
import org.apache.hadoop.yarn.server.resourcemanager.reservation.ReservationSystemUtil;
import org.apache.hadoop.yarn.server.resourcemanager.rmapp.RMApp;
import org.apache.hadoop.yarn.server.resourcemanager.rmapp.RMAppState;
import org.apache.hadoop.yarn.server.resourcemanager.rmapp.attempt.AggregateAppResourceUsage;
import org.apache.hadoop.yarn.server.resourcemanager.rmapp.attempt.RMAppAttempt;
import org.apache.hadoop.yarn.server.resourcemanager.rmapp.attempt.RMAppAttemptEvent;
import org.apache.hadoop.yarn.server.resourcemanager.rmapp.attempt.RMAppAttemptMetrics;
import org.apache.hadoop.yarn.server.resourcemanager.rmapp.attempt.RMAppAttemptState;
import org.apache.hadoop.yarn.server.resourcemanager.security.AMRMTokenSecretManager;
import org.apache.hadoop.yarn.server.resourcemanager.security.ClientToAMTokenSecretManagerInRM;
import org.apache.hadoop.yarn.server.security.MasterKeyData;
import org.apache.hadoop.yarn.util.ConverterUtils;
import org.apache.hadoop.yarn.util.resource.DefaultResourceCalculator;
import org.apache.hadoop.yarn.util.resource.ResourceCalculator;
import org.junit.Assert;

public class RMStateStoreTestBase {

  public static final Log LOG = LogFactory.getLog(RMStateStoreTestBase.class);

  protected final long epoch = 10L;

  static class TestDispatcher implements Dispatcher, EventHandler<Event> {

    ApplicationAttemptId attemptId;

    boolean notified = false;

    @SuppressWarnings("rawtypes")
    @Override
    public void register(Class<? extends Enum> eventType,
                         EventHandler handler) {
    }

    @Override
    public void handle(Event event) {
      if (event instanceof RMAppAttemptEvent) {
        RMAppAttemptEvent rmAppAttemptEvent = (RMAppAttemptEvent) event;
        assertEquals(attemptId, rmAppAttemptEvent.getApplicationAttemptId());
      }
      notified = true;
      synchronized (this) {
        notifyAll();
      }
    }

    @SuppressWarnings("rawtypes")
    @Override
    public EventHandler<Event> getEventHandler() {
      return this;
    }

  }

  public static class StoreStateVerifier {
    void afterStoreApp(RMStateStore store, ApplicationId appId) {}
    void afterStoreAppAttempt(RMStateStore store, ApplicationAttemptId
            appAttId) {}
  }

  interface RMStateStoreHelper {
    RMStateStore getRMStateStore() throws Exception;
    boolean isFinalStateValid() throws Exception;
    void writeVersion(Version version) throws Exception;
    Version getCurrentVersion() throws Exception;
    boolean appExists(RMApp app) throws Exception;
    boolean attemptExists(RMAppAttempt attempt) throws Exception;
  }

  void waitNotify(TestDispatcher dispatcher) {
    long startTime = System.currentTimeMillis();
    while(!dispatcher.notified) {
      synchronized (dispatcher) {
        try {
          dispatcher.wait(1000);
        } catch (InterruptedException e) {
          e.printStackTrace();
        }
      }
      if(System.currentTimeMillis() - startTime > 1000*60) {
        fail("Timed out attempt store notification");
      }
    }
    dispatcher.notified = false;
  }

  protected RMApp storeApp(RMStateStore store, ApplicationId appId,
      long submitTime, long startTime) throws Exception {
    ApplicationSubmissionContext context =
        new ApplicationSubmissionContextPBImpl();
    context.setApplicationId(appId);
    context.setAMContainerSpec(new ContainerLaunchContextPBImpl());

    RMApp mockApp = mock(RMApp.class);
    when(mockApp.getApplicationId()).thenReturn(appId);
    when(mockApp.getSubmitTime()).thenReturn(submitTime);
    when(mockApp.getStartTime()).thenReturn(startTime);
    when(mockApp.getApplicationSubmissionContext()).thenReturn(context);
    when(mockApp.getUser()).thenReturn("test");
    when(mockApp.getCallerContext())
        .thenReturn(new CallerContext.Builder("context").build());
    store.storeNewApplication(mockApp);
    return mockApp;
  }

  protected RMAppAttempt storeAttempt(RMStateStore store,
      ApplicationAttemptId attemptId,
      String containerIdStr, Token<AMRMTokenIdentifier> appToken,
      SecretKey clientTokenMasterKey, TestDispatcher dispatcher)
      throws Exception {

    RMAppAttemptMetrics mockRmAppAttemptMetrics = 
        mock(RMAppAttemptMetrics.class);
    Container container = new ContainerPBImpl();
    container.setId(ContainerId.fromString(containerIdStr));
    RMAppAttempt mockAttempt = mock(RMAppAttempt.class);
    when(mockAttempt.getAppAttemptId()).thenReturn(attemptId);
    when(mockAttempt.getMasterContainer()).thenReturn(container);
    when(mockAttempt.getAMRMToken()).thenReturn(appToken);
    when(mockAttempt.getClientTokenMasterKey())
        .thenReturn(clientTokenMasterKey);
    when(mockAttempt.getRMAppAttemptMetrics())
        .thenReturn(mockRmAppAttemptMetrics);
    when(mockRmAppAttemptMetrics.getAggregateAppResourceUsage())
        .thenReturn(new AggregateAppResourceUsage(new HashMap<>()));
    dispatcher.attemptId = attemptId;
    store.storeNewApplicationAttempt(mockAttempt);
    waitNotify(dispatcher);
    return mockAttempt;
  }

  protected void updateAttempt(RMStateStore store, TestDispatcher dispatcher,
      ApplicationAttemptStateData attemptState) {
    dispatcher.attemptId = attemptState.getAttemptId();
    store.updateApplicationAttemptState(attemptState);
    waitNotify(dispatcher);
  }

  void testRMAppStateStore(RMStateStoreHelper stateStoreHelper)
          throws Exception {
    testRMAppStateStore(stateStoreHelper, new StoreStateVerifier());
  }

  void testRMAppStateStore(RMStateStoreHelper stateStoreHelper,
                           StoreStateVerifier verifier)
      throws Exception {
    long submitTime = System.currentTimeMillis();
    long startTime = System.currentTimeMillis() + 1234;
    Configuration conf = new YarnConfiguration();
    RMStateStore store = stateStoreHelper.getRMStateStore();
    TestDispatcher dispatcher = new TestDispatcher();
    store.setRMDispatcher(dispatcher);

    RMContext rmContext = mock(RMContext.class);
    when(rmContext.getStateStore()).thenReturn(store);

    AMRMTokenSecretManager appTokenMgr =
        spy(new AMRMTokenSecretManager(conf, rmContext));

    MasterKeyData masterKeyData = appTokenMgr.createNewMasterKey();
    when(appTokenMgr.getMasterKey()).thenReturn(masterKeyData);

    ClientToAMTokenSecretManagerInRM clientToAMTokenMgr =
        new ClientToAMTokenSecretManagerInRM();

    ApplicationAttemptId attemptId1 = ApplicationAttemptId.fromString(
        "appattempt_1352994193343_0001_000001");
    ApplicationId appId1 = attemptId1.getApplicationId();
    storeApp(store, appId1, submitTime, startTime);
    verifier.afterStoreApp(store, appId1);

    // create application token and client token key for attempt1
    Token<AMRMTokenIdentifier> appAttemptToken1 =
        generateAMRMToken(attemptId1, appTokenMgr);
    SecretKey clientTokenKey1 =
        clientToAMTokenMgr.createMasterKey(attemptId1);

    ContainerId containerId1 = storeAttempt(store, attemptId1,
        "container_1352994193343_0001_01_000001",
        appAttemptToken1, clientTokenKey1, dispatcher)
        .getMasterContainer().getId();

    String appAttemptIdStr2 = "appattempt_1352994193343_0001_000002";
    ApplicationAttemptId attemptId2 = ApplicationAttemptId.fromString(
        appAttemptIdStr2);

    // create application token and client token key for attempt2
    Token<AMRMTokenIdentifier> appAttemptToken2 =
        generateAMRMToken(attemptId2, appTokenMgr);
    SecretKey clientTokenKey2 =
        clientToAMTokenMgr.createMasterKey(attemptId2);

    ContainerId containerId2 = storeAttempt(store, attemptId2,
        "container_1352994193343_0001_02_000001",
        appAttemptToken2, clientTokenKey2, dispatcher)
        .getMasterContainer().getId();

    ApplicationAttemptId attemptIdRemoved = ApplicationAttemptId.fromString(
        "appattempt_1352994193343_0002_000001");
    ApplicationId appIdRemoved = attemptIdRemoved.getApplicationId();
    storeApp(store, appIdRemoved, submitTime, startTime);
    storeAttempt(store, attemptIdRemoved,
        "container_1352994193343_0002_01_000001", null, null, dispatcher);
    verifier.afterStoreAppAttempt(store, attemptIdRemoved);

    RMApp mockRemovedApp = mock(RMApp.class);
    RMAppAttemptMetrics mockRmAppAttemptMetrics = 
        mock(RMAppAttemptMetrics.class);
    HashMap<ApplicationAttemptId, RMAppAttempt> attempts =
                              new HashMap<ApplicationAttemptId, RMAppAttempt>();
    ApplicationSubmissionContext context =
        new ApplicationSubmissionContextPBImpl();
    context.setApplicationId(appIdRemoved);
    when(mockRemovedApp.getSubmitTime()).thenReturn(submitTime);
    when(mockRemovedApp.getApplicationSubmissionContext()).thenReturn(context);
    when(mockRemovedApp.getAppAttempts()).thenReturn(attempts);
    when(mockRemovedApp.getUser()).thenReturn("user1");
    RMAppAttempt mockRemovedAttempt = mock(RMAppAttempt.class);
    when(mockRemovedAttempt.getAppAttemptId()).thenReturn(attemptIdRemoved);
    when(mockRemovedAttempt.getRMAppAttemptMetrics())
        .thenReturn(mockRmAppAttemptMetrics);
    when(mockRmAppAttemptMetrics.getAggregateAppResourceUsage())
        .thenReturn(new AggregateAppResourceUsage(new HashMap<>()));
    attempts.put(attemptIdRemoved, mockRemovedAttempt);
    store.removeApplication(mockRemovedApp);

    // remove application directory recursively.
    storeApp(store, appIdRemoved, submitTime, startTime);
    storeAttempt(store, attemptIdRemoved,
        "container_1352994193343_0002_01_000001", null, null, dispatcher);
    store.removeApplication(mockRemovedApp);

    // let things settle down
    Thread.sleep(1000);
    store.close();

    // give tester a chance to modify app state in the store
    modifyAppState();

    // load state
    store = stateStoreHelper.getRMStateStore();
    store.setRMDispatcher(dispatcher);
    RMState state = store.loadState();
    Map<ApplicationId, ApplicationStateData> rmAppState =
        state.getApplicationState();

    ApplicationStateData appState = rmAppState.get(appId1);
    // app is loaded
    assertNotNull(appState);
    // app is loaded correctly
    assertEquals(submitTime, appState.getSubmitTime());
    assertEquals(startTime, appState.getStartTime());
    // submission context is loaded correctly
    assertEquals(appId1,
                 appState.getApplicationSubmissionContext().getApplicationId());
    ApplicationAttemptStateData attemptState = appState.getAttempt(attemptId1);
    // attempt1 is loaded correctly
    assertNotNull(attemptState);
    assertEquals(attemptId1, attemptState.getAttemptId());
    assertEquals(-1000, attemptState.getAMContainerExitStatus());
    // attempt1 container is loaded correctly
    assertEquals(containerId1, attemptState.getMasterContainer().getId());
    // attempt1 client token master key is loaded correctly
    assertArrayEquals(
        clientTokenKey1.getEncoded(),
        attemptState.getAppAttemptTokens()
            .getSecretKey(RMStateStore.AM_CLIENT_TOKEN_MASTER_KEY_NAME));
    assertEquals("context", appState.getCallerContext().getContext());

    attemptState = appState.getAttempt(attemptId2);
    // attempt2 is loaded correctly
    assertNotNull(attemptState);
    assertEquals(attemptId2, attemptState.getAttemptId());
    // attempt2 container is loaded correctly
    assertEquals(containerId2, attemptState.getMasterContainer().getId());
    // attempt2 client token master key is loaded correctly
    assertArrayEquals(
        clientTokenKey2.getEncoded(),
        attemptState.getAppAttemptTokens()
            .getSecretKey(RMStateStore.AM_CLIENT_TOKEN_MASTER_KEY_NAME));

    //******* update application/attempt state *******//
    ApplicationStateData appState2 =
        ApplicationStateData.newInstance(appState.getSubmitTime(),
            appState.getStartTime(), appState.getUser(),
            appState.getApplicationSubmissionContext(), RMAppState.FINISHED,
            "appDiagnostics", 123, 1234, appState.getCallerContext());
    appState2.attempts.putAll(appState.attempts);
    store.updateApplicationState(appState2);

    ApplicationAttemptStateData oldAttemptState = attemptState;
    ApplicationAttemptStateData newAttemptState =
        ApplicationAttemptStateData.newInstance(
            oldAttemptState.getAttemptId(),
            oldAttemptState.getMasterContainer(),
            oldAttemptState.getAppAttemptTokens(),
            oldAttemptState.getStartTime(), RMAppAttemptState.FINISHED,
            "myTrackingUrl", "attemptDiagnostics",
            FinalApplicationStatus.SUCCEEDED, 100,
            oldAttemptState.getFinishTime(), new HashMap<>(), new HashMap<>());
    store.updateApplicationAttemptState(newAttemptState);

    // test updating the state of an app/attempt whose initial state was not
    // saved.
    ApplicationId dummyAppId = ApplicationId.newInstance(1234, 10);
    ApplicationSubmissionContext dummyContext =
        new ApplicationSubmissionContextPBImpl();
    dummyContext.setApplicationId(dummyAppId);
    dummyContext.setAMContainerSpec(new ContainerLaunchContextPBImpl());
    ApplicationStateData dummyApp =
        ApplicationStateData.newInstance(appState.getSubmitTime(),
            appState.getStartTime(), appState.getUser(), dummyContext,
            RMAppState.FINISHED, "appDiagnostics", 123, 1234, null);
    store.updateApplicationState(dummyApp);

    ApplicationAttemptId dummyAttemptId =
        ApplicationAttemptId.newInstance(dummyAppId, 6);
    ApplicationAttemptStateData dummyAttempt =
        ApplicationAttemptStateData.newInstance(dummyAttemptId,
            oldAttemptState.getMasterContainer(),
            oldAttemptState.getAppAttemptTokens(),
            oldAttemptState.getStartTime(), RMAppAttemptState.FINISHED,
            "myTrackingUrl", "attemptDiagnostics",
            FinalApplicationStatus.SUCCEEDED, 111,
            oldAttemptState.getFinishTime(), new HashMap<>(), new HashMap<>());
    store.updateApplicationAttemptState(dummyAttempt);

    // let things settle down
    Thread.sleep(1000);
    store.close();

    // check updated application state.
    store = stateStoreHelper.getRMStateStore();
    store.setRMDispatcher(dispatcher);
    RMState newRMState = store.loadState();
    Map<ApplicationId, ApplicationStateData> newRMAppState =
        newRMState.getApplicationState();
    assertNotNull(newRMAppState.get(
        dummyApp.getApplicationSubmissionContext().getApplicationId()));
    ApplicationStateData updatedAppState = newRMAppState.get(appId1);
    assertEquals(appState.getApplicationSubmissionContext().getApplicationId(),
        updatedAppState.getApplicationSubmissionContext().getApplicationId());
    assertEquals(appState.getSubmitTime(), updatedAppState.getSubmitTime());
    assertEquals(appState.getStartTime(), updatedAppState.getStartTime());
    assertEquals(appState.getUser(), updatedAppState.getUser());
    // new app state fields
    assertEquals( RMAppState.FINISHED, updatedAppState.getState());
    assertEquals("appDiagnostics", updatedAppState.getDiagnostics());
    assertEquals(1234, updatedAppState.getFinishTime());

    // check updated attempt state
    assertNotNull(newRMAppState.get(dummyApp.getApplicationSubmissionContext
        ().getApplicationId()).getAttempt(dummyAttemptId));
    ApplicationAttemptStateData updatedAttemptState =
        updatedAppState.getAttempt(newAttemptState.getAttemptId());
    assertEquals(oldAttemptState.getAttemptId(),
      updatedAttemptState.getAttemptId());
    assertEquals(containerId2, updatedAttemptState.getMasterContainer().getId());
    assertArrayEquals(
        clientTokenKey2.getEncoded(),
        attemptState.getAppAttemptTokens()
            .getSecretKey(RMStateStore.AM_CLIENT_TOKEN_MASTER_KEY_NAME));
    // new attempt state fields
    assertEquals(RMAppAttemptState.FINISHED, updatedAttemptState.getState());
    assertEquals("myTrackingUrl", updatedAttemptState.getFinalTrackingUrl());
    assertEquals("attemptDiagnostics", updatedAttemptState.getDiagnostics());
    assertEquals(100, updatedAttemptState.getAMContainerExitStatus());
    assertEquals(FinalApplicationStatus.SUCCEEDED,
      updatedAttemptState.getFinalApplicationStatus());

    // assert store is in expected state after everything is cleaned
    assertTrue(stateStoreHelper.isFinalStateValid());

    store.close();
  }

  public void testRMDTSecretManagerStateStore(
      RMStateStoreHelper stateStoreHelper) throws Exception {
    RMStateStore store = stateStoreHelper.getRMStateStore();
    TestDispatcher dispatcher = new TestDispatcher();
    store.setRMDispatcher(dispatcher);

    // store RM delegation token;
    RMDelegationTokenIdentifier dtId1 =
        new RMDelegationTokenIdentifier(new Text("owner1"),
          new Text("renewer1"), new Text("realuser1"));
    int sequenceNumber = 1111;
    dtId1.setSequenceNumber(sequenceNumber);
    byte[] tokenBeforeStore = dtId1.getBytes();
    Long renewDate1 = new Long(System.currentTimeMillis());
    store.storeRMDelegationToken(dtId1, renewDate1);
    modifyRMDelegationTokenState();
    Map<RMDelegationTokenIdentifier, Long> token1 =
        new HashMap<RMDelegationTokenIdentifier, Long>();
    token1.put(dtId1, renewDate1);
    // store delegation key;
    DelegationKey key = new DelegationKey(1234, 4321 , "keyBytes".getBytes());
    HashSet<DelegationKey> keySet = new HashSet<DelegationKey>();
    keySet.add(key);
    store.storeRMDTMasterKey(key);

    RMDTSecretManagerState secretManagerState =
        store.loadState().getRMDTSecretManagerState();
    Assert.assertEquals(token1, secretManagerState.getTokenState());
    Assert.assertEquals(keySet, secretManagerState.getMasterKeyState());
    Assert.assertEquals(sequenceNumber,
        secretManagerState.getDTSequenceNumber());
    RMDelegationTokenIdentifier tokenAfterStore =
        secretManagerState.getTokenState().keySet().iterator().next();
    Assert.assertTrue(Arrays.equals(tokenBeforeStore,
      tokenAfterStore.getBytes()));

    // update RM delegation token;
    renewDate1 = new Long(System.currentTimeMillis());
    store.updateRMDelegationToken(dtId1, renewDate1);
    token1.put(dtId1, renewDate1);

    RMDTSecretManagerState updateSecretManagerState =
        store.loadState().getRMDTSecretManagerState();
    Assert.assertEquals(token1, updateSecretManagerState.getTokenState());
    Assert.assertEquals(keySet, updateSecretManagerState.getMasterKeyState());
    Assert.assertEquals(sequenceNumber,
        updateSecretManagerState.getDTSequenceNumber());

    // check to delete delegationKey
    store.removeRMDTMasterKey(key);
    keySet.clear();
    RMDTSecretManagerState noKeySecretManagerState =
        store.loadState().getRMDTSecretManagerState();
    Assert.assertEquals(token1, noKeySecretManagerState.getTokenState());
    Assert.assertEquals(keySet, noKeySecretManagerState.getMasterKeyState());
    Assert.assertEquals(sequenceNumber,
        noKeySecretManagerState.getDTSequenceNumber());

    // check to delete delegationToken
    store.removeRMDelegationToken(dtId1);
    RMDTSecretManagerState noKeyAndTokenSecretManagerState =
        store.loadState().getRMDTSecretManagerState();
    token1.clear();
    Assert.assertEquals(token1,
        noKeyAndTokenSecretManagerState.getTokenState());
    Assert.assertEquals(keySet,
        noKeyAndTokenSecretManagerState.getMasterKeyState());
    Assert.assertEquals(sequenceNumber,
        noKeySecretManagerState.getDTSequenceNumber());
    store.close();

  }

  protected Token<AMRMTokenIdentifier> generateAMRMToken(
      ApplicationAttemptId attemptId,
      AMRMTokenSecretManager appTokenMgr) {
    Token<AMRMTokenIdentifier> appToken =
        appTokenMgr.createAndGetAMRMToken(attemptId);
    appToken.setService(new Text("appToken service"));
    return appToken;
  }

  public void testCheckVersion(RMStateStoreHelper stateStoreHelper)
      throws Exception {
    RMStateStore store = stateStoreHelper.getRMStateStore();
    store.setRMDispatcher(new TestDispatcher());

    // default version
    Version defaultVersion = stateStoreHelper.getCurrentVersion();
    store.checkVersion();
    Assert.assertEquals(defaultVersion, store.loadVersion());

    // compatible version
    Version compatibleVersion =
        Version.newInstance(defaultVersion.getMajorVersion(),
          defaultVersion.getMinorVersion() + 2);
    stateStoreHelper.writeVersion(compatibleVersion);
    Assert.assertEquals(compatibleVersion, store.loadVersion());
    store.checkVersion();
    // overwrite the compatible version
    Assert.assertEquals(defaultVersion, store.loadVersion());

    // incompatible version
    Version incompatibleVersion =
        Version.newInstance(defaultVersion.getMajorVersion() + 2,
          defaultVersion.getMinorVersion());
    stateStoreHelper.writeVersion(incompatibleVersion);
    try {
      store.checkVersion();
      Assert.fail("Invalid version, should fail.");
    } catch (Throwable t) {
      Assert.assertTrue(t instanceof RMStateVersionIncompatibleException);
    }
  }
  
  public void testEpoch(RMStateStoreHelper stateStoreHelper)
      throws Exception {
    RMStateStore store = stateStoreHelper.getRMStateStore();
    store.setRMDispatcher(new TestDispatcher());
    
    long firstTimeEpoch = store.getAndIncrementEpoch();
    Assert.assertEquals(epoch, firstTimeEpoch);
    
    long secondTimeEpoch = store.getAndIncrementEpoch();
    Assert.assertEquals(epoch + 1, secondTimeEpoch);
    
    long thirdTimeEpoch = store.getAndIncrementEpoch();
    Assert.assertEquals(epoch + 2, thirdTimeEpoch);
  }

  public void testAppDeletion(RMStateStoreHelper stateStoreHelper)
      throws Exception {
    RMStateStore store = stateStoreHelper.getRMStateStore();
    store.setRMDispatcher(new TestDispatcher());
    ArrayList<RMApp> appList = createAndStoreApps(stateStoreHelper, store, 5);

    for (RMApp app : appList) {
      // remove the app
      store.removeApplication(app);
      // wait for app to be removed.
      while (true) {
        if (!stateStoreHelper.appExists(app)) {
          break;
        } else {
          Thread.sleep(100);
        }
      }
    }
  }

  private ArrayList<RMApp> createAndStoreApps(
      RMStateStoreHelper stateStoreHelper, RMStateStore store, int numApps)
      throws Exception {
    ArrayList<RMApp> appList = new ArrayList<RMApp>();
    for (int i = 0; i < numApps; i++) {
      ApplicationId appId = ApplicationId.newInstance(1383183338, i);
      RMApp app = storeApp(store, appId, 123456789, 987654321);
      appList.add(app);
    }

    Assert.assertEquals(numApps, appList.size());
    for (RMApp app : appList) {
      // wait for app to be stored.
      while (true) {
        if (stateStoreHelper.appExists(app)) {
          break;
        } else {
          Thread.sleep(100);
        }
      }
    }
    return appList;
  }

  public void testDeleteStore(RMStateStoreHelper stateStoreHelper)
      throws Exception {
    RMStateStore store = stateStoreHelper.getRMStateStore();
    ArrayList<RMApp> appList = createAndStoreApps(stateStoreHelper, store, 5);
    store.deleteStore();
    // verify apps deleted
    for (RMApp app : appList) {
      Assert.assertFalse(stateStoreHelper.appExists(app));
    }
  }

  public void testRemoveApplication(RMStateStoreHelper stateStoreHelper)
      throws Exception {
    RMStateStore store = stateStoreHelper.getRMStateStore();
    int noOfApps = 2;
    ArrayList<RMApp> appList =
        createAndStoreApps(stateStoreHelper, store, noOfApps);

    RMApp rmApp1 = appList.get(0);
    store.removeApplication(rmApp1.getApplicationId());
    Assert.assertFalse(stateStoreHelper.appExists(rmApp1));

    RMApp rmApp2 = appList.get(1);
    Assert.assertTrue(stateStoreHelper.appExists(rmApp2));
  }

  public void testRemoveAttempt(RMStateStoreHelper stateStoreHelper)
    throws Exception {
    RMStateStore store = stateStoreHelper.getRMStateStore();
    TestDispatcher dispatcher = new TestDispatcher();
    store.setRMDispatcher(dispatcher);

    ApplicationId appId = ApplicationId.newInstance(1383183339, 6);
    storeApp(store, appId, 123456, 564321);

    ApplicationAttemptId attemptId1 =
        ApplicationAttemptId.newInstance(appId, 1);
    RMAppAttempt attempt1 = storeAttempt(store, attemptId1,
        ContainerId.newContainerId(attemptId1, 1).toString(),
        null, null, dispatcher);
    ApplicationAttemptId attemptId2 =
        ApplicationAttemptId.newInstance(appId, 2);
    RMAppAttempt attempt2 = storeAttempt(store, attemptId2,
        ContainerId.newContainerId(attemptId2, 1).toString(),
        null, null, dispatcher);
    store.removeApplicationAttemptInternal(attemptId1);
    Assert.assertFalse(stateStoreHelper.attemptExists(attempt1));
    Assert.assertTrue(stateStoreHelper.attemptExists(attempt2));

    // let things settle down
    Thread.sleep(1000);
    store.close();

    // load state
    store = stateStoreHelper.getRMStateStore();
    RMState state = store.loadState();
    Map<ApplicationId, ApplicationStateData> rmAppState =
        state.getApplicationState();

    ApplicationStateData appState = rmAppState.get(appId);
    // app is loaded
    assertNotNull(appState);
    assertEquals(2, appState.getFirstAttemptId());
    assertNull(appState.getAttempt(attemptId1));
    assertNotNull(appState.getAttempt(attemptId2));
  }

  protected void modifyAppState() throws Exception {

  }

  protected void modifyRMDelegationTokenState() throws Exception {

  }

  public void testAMRMTokenSecretManagerStateStore(
      RMStateStoreHelper stateStoreHelper) throws Exception {
    System.out.println("Start testing");
    RMStateStore store = stateStoreHelper.getRMStateStore();
    TestDispatcher dispatcher = new TestDispatcher();
    store.setRMDispatcher(dispatcher);

    RMContext rmContext = mock(RMContext.class);
    when(rmContext.getStateStore()).thenReturn(store);
    Configuration conf = new YarnConfiguration();
    AMRMTokenSecretManager appTokenMgr =
        new AMRMTokenSecretManager(conf, rmContext);

    //create and save the first masterkey
    MasterKeyData firstMasterKeyData = appTokenMgr.createNewMasterKey();

    AMRMTokenSecretManagerState state1 =
        AMRMTokenSecretManagerState.newInstance(
          firstMasterKeyData.getMasterKey(), null);
    rmContext.getStateStore()
        .storeOrUpdateAMRMTokenSecretManager(state1,
      false);

    // load state
    store = stateStoreHelper.getRMStateStore();
    when(rmContext.getStateStore()).thenReturn(store);
    store.setRMDispatcher(dispatcher);
    RMState state = store.loadState();
    Assert.assertNotNull(state.getAMRMTokenSecretManagerState());
    Assert.assertEquals(firstMasterKeyData.getMasterKey(), state
      .getAMRMTokenSecretManagerState().getCurrentMasterKey());
    Assert.assertNull(state
      .getAMRMTokenSecretManagerState().getNextMasterKey());

    //create and save the second masterkey
    MasterKeyData secondMasterKeyData = appTokenMgr.createNewMasterKey();
    AMRMTokenSecretManagerState state2 =
        AMRMTokenSecretManagerState
          .newInstance(firstMasterKeyData.getMasterKey(),
            secondMasterKeyData.getMasterKey());
    rmContext.getStateStore().storeOrUpdateAMRMTokenSecretManager(state2,
      true);

    // load state
    store = stateStoreHelper.getRMStateStore();
    when(rmContext.getStateStore()).thenReturn(store);
    store.setRMDispatcher(dispatcher);
    RMState state_2 = store.loadState();
    Assert.assertNotNull(state_2.getAMRMTokenSecretManagerState());
    Assert.assertEquals(firstMasterKeyData.getMasterKey(), state_2
      .getAMRMTokenSecretManagerState().getCurrentMasterKey());
    Assert.assertEquals(secondMasterKeyData.getMasterKey(), state_2
      .getAMRMTokenSecretManagerState().getNextMasterKey());

    // re-create the masterKeyData based on the recovered masterkey
    // should have the same secretKey
    appTokenMgr.recover(state_2);
    Assert.assertEquals(appTokenMgr.getCurrnetMasterKeyData().getSecretKey(),
      firstMasterKeyData.getSecretKey());
    Assert.assertEquals(appTokenMgr.getNextMasterKeyData().getSecretKey(),
      secondMasterKeyData.getSecretKey());

    store.close();
  }

  public void testReservationStateStore(
      RMStateStoreHelper stateStoreHelper) throws Exception {
    RMStateStore store = stateStoreHelper.getRMStateStore();
    TestDispatcher dispatcher = new TestDispatcher();
    store.setRMDispatcher(dispatcher);

    RMContext rmContext = mock(RMContext.class);
    when(rmContext.getStateStore()).thenReturn(store);

    long ts = System.currentTimeMillis();
    ReservationId r1 = ReservationId.newInstance(ts, 1);
    int start = 1;
    int[] alloc = { 10, 10, 10, 10, 10 };
    ResourceCalculator res = new DefaultResourceCalculator();
    Resource minAlloc = Resource.newInstance(1024, 1);
    boolean hasGang = true;
    String planName = "dedicated";
    ReservationDefinition rDef =
        ReservationSystemTestUtil.createSimpleReservationDefinition(
            start, start + alloc.length + 1, alloc.length);
    ReservationAllocation allocation = new InMemoryReservationAllocation(
        r1, rDef, "u3", planName, 0, 0 + alloc.length,
        ReservationSystemTestUtil.generateAllocation(0L, 1L, alloc), res,
        minAlloc, hasGang);
    ReservationAllocationStateProto allocationStateProto =
        ReservationSystemUtil.buildStateProto(allocation);
    assertAllocationStateEqual(allocation, allocationStateProto);

    // 1. Load empty store and verify no errors
    store = stateStoreHelper.getRMStateStore();
    when(rmContext.getStateStore()).thenReturn(store);
    store.setRMDispatcher(dispatcher);
    RMState state = store.loadState();
    Map<String, Map<ReservationId, ReservationAllocationStateProto>>
      reservationState = state.getReservationState();
    Assert.assertNotNull(reservationState);

    // 2. Store single reservation and verify
    String reservationIdName = r1.toString();
    rmContext.getStateStore().storeNewReservation(
        allocationStateProto,
        planName, reservationIdName);


    // load state and verify new state
    validateStoredReservation(
        stateStoreHelper, dispatcher, rmContext, r1, planName, allocation,
        allocationStateProto);

    // 3. update state test
    alloc = new int[]{6, 6, 6};
    hasGang = false;
    allocation = new InMemoryReservationAllocation(
        r1, rDef, "u3", planName, 2, 2 + alloc.length,
        ReservationSystemTestUtil.generateAllocation(1L, 2L, alloc), res,
        minAlloc, hasGang);
    allocationStateProto =
        ReservationSystemUtil.buildStateProto(allocation);
    rmContext.getStateStore().removeReservation(planName, reservationIdName);
    rmContext.getStateStore().storeNewReservation(allocationStateProto, planName, reservationIdName);

    // load state and verify updated reservation
    validateStoredReservation(
        stateStoreHelper, dispatcher, rmContext, r1, planName, allocation,
        allocationStateProto);

    // 4. add a second one and remove the first one
    ReservationId r2 = ReservationId.newInstance(ts, 2);
    ReservationAllocation allocation2 = new InMemoryReservationAllocation(
        r2, rDef, "u3", planName, 0, 0 + alloc.length,
        ReservationSystemTestUtil.generateAllocation(0L, 1L, alloc), res,
        minAlloc, hasGang);
    ReservationAllocationStateProto allocationStateProto2 =
        ReservationSystemUtil.buildStateProto(allocation2);
    String reservationIdName2 = r2.toString();

    rmContext.getStateStore().storeNewReservation(
        allocationStateProto2,
        planName, reservationIdName2);
    rmContext.getStateStore().removeReservation(planName, reservationIdName);

    // load state and verify r1 is removed and r2 is still there
    Map<ReservationId, ReservationAllocationStateProto> reservations;

    store = stateStoreHelper.getRMStateStore();
    when(rmContext.getStateStore()).thenReturn(store);
    store.setRMDispatcher(dispatcher);
    state = store.loadState();
    reservationState = state.getReservationState();
    Assert.assertNotNull(reservationState);
    reservations = reservationState.get(planName);
    Assert.assertNotNull(reservations);
    ReservationAllocationStateProto storedReservationAllocation =
        reservations.get(r1);
    Assert.assertNull("Removed reservation should not be available in store",
        storedReservationAllocation);

    storedReservationAllocation = reservations.get(r2);
    assertAllocationStateEqual(
        allocationStateProto2, storedReservationAllocation);
    assertAllocationStateEqual(allocation2, storedReservationAllocation);


    // 5. remove last reservation removes the plan state
    rmContext.getStateStore().removeReservation(planName, reservationIdName2);

    store = stateStoreHelper.getRMStateStore();
    when(rmContext.getStateStore()).thenReturn(store);
    store.setRMDispatcher(dispatcher);
    state = store.loadState();
    reservationState = state.getReservationState();
    Assert.assertNotNull(reservationState);
    reservations = reservationState.get(planName);
    Assert.assertNull(reservations);
  }

  private void validateStoredReservation(
      RMStateStoreHelper stateStoreHelper, TestDispatcher dispatcher,
      RMContext rmContext, ReservationId r1, String planName,
      ReservationAllocation allocation,
      ReservationAllocationStateProto allocationStateProto) throws Exception {
    RMStateStore store = stateStoreHelper.getRMStateStore();
    when(rmContext.getStateStore()).thenReturn(store);
    store.setRMDispatcher(dispatcher);
    RMState state = store.loadState();
    Map<String, Map<ReservationId, ReservationAllocationStateProto>>
        reservationState = state.getReservationState();
    Assert.assertNotNull(reservationState);
    Map<ReservationId, ReservationAllocationStateProto> reservations =
        reservationState.get(planName);
    Assert.assertNotNull(reservations);
    ReservationAllocationStateProto storedReservationAllocation =
        reservations.get(r1);
    Assert.assertNotNull(storedReservationAllocation);

    assertAllocationStateEqual(
        allocationStateProto, storedReservationAllocation);
    assertAllocationStateEqual(allocation, storedReservationAllocation);
  }

  void assertAllocationStateEqual(
      ReservationAllocationStateProto expected,
      ReservationAllocationStateProto actual) {

    Assert.assertEquals(
        expected.getAcceptanceTime(), actual.getAcceptanceTime());
    Assert.assertEquals(expected.getStartTime(), actual.getStartTime());
    Assert.assertEquals(expected.getEndTime(), actual.getEndTime());
    Assert.assertEquals(expected.getContainsGangs(), actual.getContainsGangs());
    Assert.assertEquals(expected.getUser(), actual.getUser());
    assertEquals(
        expected.getReservationDefinition(), actual.getReservationDefinition());
    assertEquals(expected.getAllocationRequestsList(),
        actual.getAllocationRequestsList());
  }

  void assertAllocationStateEqual(
      ReservationAllocation expected,
      ReservationAllocationStateProto actual) {
    Assert.assertEquals(
        expected.getAcceptanceTime(), actual.getAcceptanceTime());
    Assert.assertEquals(expected.getStartTime(), actual.getStartTime());
    Assert.assertEquals(expected.getEndTime(), actual.getEndTime());
    Assert.assertEquals(expected.containsGangs(), actual.getContainsGangs());
    Assert.assertEquals(expected.getUser(), actual.getUser());
    assertEquals(
        expected.getReservationDefinition(),
        ReservationSystemUtil.convertFromProtoFormat(
            actual.getReservationDefinition()));
    assertEquals(
        expected.getAllocationRequests(),
        ReservationSystemUtil.toAllocations(
            actual.getAllocationRequestsList()));
  }
}