/**
 * 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
 * <p>
 * http://www.apache.org/licenses/LICENSE-2.0
 * <p>
 * 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.federation;

import java.io.IOException;
import java.io.StringReader;
import java.net.UnknownHostException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import javax.xml.bind.JAXBException;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.ha.HAServiceProtocol;
import org.apache.hadoop.test.GenericTestUtils;
import org.apache.hadoop.test.LambdaTestUtils;
import org.apache.hadoop.logging.LogCapturer;
import org.apache.hadoop.util.Time;
import org.apache.hadoop.yarn.MockApps;
import org.apache.hadoop.yarn.api.records.ApplicationId;
import org.apache.hadoop.yarn.api.records.ApplicationSubmissionContext;
import org.apache.hadoop.yarn.api.records.Priority;
import org.apache.hadoop.yarn.api.records.impl.pb.ApplicationSubmissionContextPBImpl;
import org.apache.hadoop.yarn.conf.YarnConfiguration;
import org.apache.hadoop.yarn.exceptions.YarnException;
import org.apache.hadoop.yarn.server.federation.policies.exceptions.FederationPolicyInitializationException;
import org.apache.hadoop.yarn.server.federation.policies.manager.UniformBroadcastPolicyManager;
import org.apache.hadoop.yarn.server.federation.store.FederationStateStore;
import org.apache.hadoop.yarn.server.federation.store.exception.FederationStateStoreException;
import org.apache.hadoop.yarn.server.federation.store.records.GetSubClusterInfoRequest;
import org.apache.hadoop.yarn.server.federation.store.records.GetSubClusterInfoResponse;
import org.apache.hadoop.yarn.server.federation.store.records.SubClusterDeregisterRequest;
import org.apache.hadoop.yarn.server.federation.store.records.SubClusterId;
import org.apache.hadoop.yarn.server.federation.store.records.SubClusterInfo;
import org.apache.hadoop.yarn.server.federation.store.records.SubClusterState;
import org.apache.hadoop.yarn.server.federation.store.records.ApplicationHomeSubCluster;
import org.apache.hadoop.yarn.server.federation.store.records.AddApplicationHomeSubClusterRequest;
import org.apache.hadoop.yarn.server.federation.store.records.GetApplicationHomeSubClusterRequest;
import org.apache.hadoop.yarn.server.federation.store.records.GetApplicationHomeSubClusterResponse;
import org.apache.hadoop.yarn.server.federation.store.records.GetApplicationsHomeSubClusterRequest;
import org.apache.hadoop.yarn.server.federation.store.records.GetApplicationsHomeSubClusterResponse;
import org.apache.hadoop.yarn.server.federation.store.records.SubClusterPolicyConfiguration;
import org.apache.hadoop.yarn.server.federation.store.records.SetSubClusterPolicyConfigurationRequest;
import org.apache.hadoop.yarn.server.federation.store.records.GetSubClusterPolicyConfigurationRequest;
import org.apache.hadoop.yarn.server.federation.store.records.GetSubClusterPolicyConfigurationResponse;
import org.apache.hadoop.yarn.server.federation.store.records.GetSubClusterPoliciesConfigurationsRequest;
import org.apache.hadoop.yarn.server.federation.store.records.GetSubClusterPoliciesConfigurationsResponse;
import org.apache.hadoop.yarn.server.federation.store.records.SubClusterRegisterRequest;
import org.apache.hadoop.yarn.server.federation.store.records.GetSubClustersInfoRequest;
import org.apache.hadoop.yarn.server.federation.store.records.GetSubClustersInfoResponse;
import org.apache.hadoop.yarn.server.federation.store.records.SubClusterHeartbeatRequest;
import org.apache.hadoop.yarn.server.federation.store.records.SubClusterHeartbeatResponse;
import org.apache.hadoop.yarn.server.resourcemanager.ApplicationMasterService;
import org.apache.hadoop.yarn.server.resourcemanager.MockRM;
import org.apache.hadoop.yarn.server.resourcemanager.RMAppManager;
import org.apache.hadoop.yarn.server.resourcemanager.RMContext;
import org.apache.hadoop.yarn.server.resourcemanager.recovery.MemoryRMStateStore;
import org.apache.hadoop.yarn.server.resourcemanager.rmapp.RMApp;
import org.apache.hadoop.yarn.server.resourcemanager.rmapp.RMAppImpl;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.YarnScheduler;
import org.apache.hadoop.yarn.server.resourcemanager.webapp.dao.ClusterMetricsInfo;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;

import com.sun.jersey.api.json.JSONConfiguration;
import com.sun.jersey.api.json.JSONJAXBContext;
import com.sun.jersey.api.json.JSONUnmarshaller;

import static org.mockito.Mockito.mock;

/**
 * Unit tests for FederationStateStoreService.
 */
public class TestFederationRMStateStoreService {

  private final HAServiceProtocol.StateChangeRequestInfo requestInfo =
      new HAServiceProtocol.StateChangeRequestInfo(
          HAServiceProtocol.RequestSource.REQUEST_BY_USER);
  private final SubClusterId subClusterId = SubClusterId.newInstance("SC-1");
  private final GetSubClusterInfoRequest request =
      GetSubClusterInfoRequest.newInstance(subClusterId);

  private Configuration conf;
  private FederationStateStore stateStore;
  private long lastHearbeatTS = 0;
  private JSONJAXBContext jc;
  private JSONUnmarshaller unmarshaller;
  private MockRM mockRM;

  @Before
  public void setUp() throws IOException, YarnException, JAXBException {
    conf = new YarnConfiguration();
    jc = new JSONJAXBContext(
        JSONConfiguration.mapped().rootUnwrapping(false).build(),
        ClusterMetricsInfo.class);
    unmarshaller = jc.createJSONUnmarshaller();

    conf.setBoolean(YarnConfiguration.FEDERATION_ENABLED, true);
    conf.setInt(YarnConfiguration.FEDERATION_STATESTORE_HEARTBEAT_INITIAL_DELAY, 10);
    conf.set(YarnConfiguration.RM_CLUSTER_ID, subClusterId.getId());

    // set up MockRM
    mockRM = new MockRM(conf);
    mockRM.init(conf);
    mockRM.start();
  }

  @After
  public void tearDown() throws Exception {
    unmarshaller = null;
    jc = null;
    mockRM.stop();
    mockRM = null;
  }

  @Test
  public void testFederationStateStoreService() throws Exception {
    conf.setBoolean(YarnConfiguration.FEDERATION_ENABLED, true);
    conf.set(YarnConfiguration.RM_CLUSTER_ID, subClusterId.getId());
    final MockRM rm = new MockRM(conf);

    // Initially there should be no entry for the sub-cluster
    rm.init(conf);
    stateStore = rm.getFederationStateStoreService().getStateStoreClient();
    GetSubClusterInfoResponse response = stateStore.getSubCluster(request);
    Assert.assertNull(response);

    // Validate if sub-cluster is registered
    rm.start();
    String capability = checkSubClusterInfo(SubClusterState.SC_NEW);
    Assert.assertTrue(capability.isEmpty());

    // Heartbeat to see if sub-cluster transitions to running
    FederationStateStoreHeartbeat storeHeartbeat =
        rm.getFederationStateStoreService().getStateStoreHeartbeatThread();
    storeHeartbeat.run();
    capability = checkSubClusterInfo(SubClusterState.SC_RUNNING);
    checkClusterMetricsInfo(capability, 0);

    // heartbeat again after adding a node.
    rm.registerNode("127.0.0.1:1234", 4 * 1024);
    storeHeartbeat.run();
    capability = checkSubClusterInfo(SubClusterState.SC_RUNNING);
    checkClusterMetricsInfo(capability, 1);

    // Validate sub-cluster deregistration
    rm.getFederationStateStoreService()
        .deregisterSubCluster(SubClusterDeregisterRequest
            .newInstance(subClusterId, SubClusterState.SC_UNREGISTERED));
    checkSubClusterInfo(SubClusterState.SC_UNREGISTERED);

    // check after failover
    explicitFailover(rm);

    capability = checkSubClusterInfo(SubClusterState.SC_NEW);
    Assert.assertTrue(capability.isEmpty());

    // Heartbeat to see if sub-cluster transitions to running
    storeHeartbeat =
        rm.getFederationStateStoreService().getStateStoreHeartbeatThread();
    storeHeartbeat.run();
    capability = checkSubClusterInfo(SubClusterState.SC_RUNNING);
    checkClusterMetricsInfo(capability, 0);

    // heartbeat again after adding a node.
    rm.registerNode("127.0.0.1:1234", 4 * 1024);
    storeHeartbeat.run();
    capability = checkSubClusterInfo(SubClusterState.SC_RUNNING);
    checkClusterMetricsInfo(capability, 1);

    rm.stop();
  }

  private void explicitFailover(MockRM rm) throws IOException {
    rm.getAdminService().transitionToStandby(requestInfo);
    Assert.assertTrue(rm.getRMContext()
        .getHAServiceState() == HAServiceProtocol.HAServiceState.STANDBY);
    rm.getAdminService().transitionToActive(requestInfo);
    Assert.assertTrue(rm.getRMContext()
        .getHAServiceState() == HAServiceProtocol.HAServiceState.ACTIVE);
    lastHearbeatTS = 0;
    stateStore = rm.getFederationStateStoreService().getStateStoreClient();
  }

  private void checkClusterMetricsInfo(String capability, int numNodes)
      throws JAXBException {
    ClusterMetricsInfo clusterMetricsInfo = unmarshaller.unmarshalFromJSON(
        new StringReader(capability), ClusterMetricsInfo.class);
    Assert.assertEquals(numNodes, clusterMetricsInfo.getTotalNodes());
  }

  private String checkSubClusterInfo(SubClusterState state)
      throws YarnException, UnknownHostException {
    Assert.assertNotNull(stateStore.getSubCluster(request));
    SubClusterInfo response =
        stateStore.getSubCluster(request).getSubClusterInfo();
    Assert.assertEquals(state, response.getState());
    Assert.assertTrue(response.getLastHeartBeat() >= lastHearbeatTS);
    String expectedAddress =
        (response.getClientRMServiceAddress().split(":"))[0];
    Assert.assertEquals(expectedAddress,
        (response.getAMRMServiceAddress().split(":"))[0]);
    Assert.assertEquals(expectedAddress,
        (response.getRMAdminServiceAddress().split(":"))[0]);
    Assert.assertEquals(expectedAddress,
        (response.getRMWebServiceAddress().split(":"))[0]);
    lastHearbeatTS = response.getLastHeartBeat();
    return response.getCapability();
  }

  @Test
  public void testFederationStateStoreServiceInitialHeartbeatDelay() throws Exception {
    conf.setBoolean(YarnConfiguration.FEDERATION_ENABLED, true);
    conf.setInt(YarnConfiguration.FEDERATION_STATESTORE_HEARTBEAT_INITIAL_DELAY, 10);
    conf.set(YarnConfiguration.RM_CLUSTER_ID, subClusterId.getId());

    LogCapturer logCapture =
        LogCapturer.captureLogs(FederationStateStoreService.LOG);

    final MockRM rm = new MockRM(conf);

    // Initially there should be no entry for the sub-cluster
    rm.init(conf);
    stateStore = rm.getFederationStateStoreService().getStateStoreClient();
    GetSubClusterInfoResponse response = stateStore.getSubCluster(request);
    Assert.assertNull(response);

    // Validate if sub-cluster is registered
    rm.start();
    String capability = checkSubClusterInfo(SubClusterState.SC_NEW);
    Assert.assertTrue(capability.isEmpty());

    // Heartbeat to see if sub-cluster transitions to running
    FederationStateStoreHeartbeat storeHeartbeat =
        rm.getFederationStateStoreService().getStateStoreHeartbeatThread();
    storeHeartbeat.run();
    capability = checkSubClusterInfo(SubClusterState.SC_RUNNING);
    checkClusterMetricsInfo(capability, 0);

    Assert.assertTrue(logCapture.getOutput().contains(
        "Started federation membership heartbeat with interval: 300 and initial delay: 10"));
    rm.stop();
  }

  @Test
  public void testCleanUpApplication() throws Exception {

    // set yarn configuration
    conf.setBoolean(YarnConfiguration.FEDERATION_ENABLED, true);
    conf.setInt(YarnConfiguration.FEDERATION_STATESTORE_HEARTBEAT_INITIAL_DELAY, 10);
    conf.set(YarnConfiguration.RM_CLUSTER_ID, subClusterId.getId());

    // set up MockRM
    final MockRM rm = new MockRM(conf);
    rm.init(conf);
    stateStore = rm.getFederationStateStoreService().getStateStoreClient();
    rm.start();

    // init subCluster Heartbeat,
    // and check that the subCluster is in a running state
    FederationStateStoreService stateStoreService = rm.getFederationStateStoreService();
    FederationStateStoreHeartbeat storeHeartbeat = stateStoreService.getStateStoreHeartbeatThread();
    storeHeartbeat.run();
    checkSubClusterInfo(SubClusterState.SC_RUNNING);

    // generate an application and join the [SC-1] cluster
    ApplicationId appId = ApplicationId.newInstance(Time.now(), 1);
    addApplication2StateStore(appId, stateStore);

    // make sure the app can be queried in the stateStore
    GetApplicationHomeSubClusterRequest appRequest =
         GetApplicationHomeSubClusterRequest.newInstance(appId);
    GetApplicationHomeSubClusterResponse response =
         stateStore.getApplicationHomeSubCluster(appRequest);
    Assert.assertNotNull(response);
    ApplicationHomeSubCluster appHomeSubCluster = response.getApplicationHomeSubCluster();
    Assert.assertNotNull(appHomeSubCluster);
    Assert.assertNotNull(appHomeSubCluster.getApplicationId());
    Assert.assertEquals(appId, appHomeSubCluster.getApplicationId());

    // clean up the app.
    boolean cleanUpResult =
        stateStoreService.cleanUpFinishApplicationsWithRetries(appId, true);
    Assert.assertTrue(cleanUpResult);

    // after clean, the app can no longer be queried from the stateStore.
    LambdaTestUtils.intercept(FederationStateStoreException.class,
        "Application " + appId + " does not exist",
        () -> stateStore.getApplicationHomeSubCluster(appRequest));

  }

  @Test
  public void testCleanUpApplicationWhenRMStart() throws Exception {

    // We design such a test case.
    // Step1. We add app01, app02, app03 to the stateStore,
    // But these apps are not in RM's RMContext, they are finished apps
    // Step2. We simulate RM startup, there is only app04 in RMContext.
    // Step3. We wait for 5 seconds, the automatic cleanup thread should clean up finished apps.

    // set yarn configuration.
    conf.setBoolean(YarnConfiguration.FEDERATION_ENABLED, true);
    conf.setInt(YarnConfiguration.FEDERATION_STATESTORE_HEARTBEAT_INITIAL_DELAY, 10);
    conf.set(YarnConfiguration.RM_CLUSTER_ID, subClusterId.getId());
    conf.setBoolean(YarnConfiguration.RECOVERY_ENABLED, true);

    // set up MockRM.
    MockRM rm = new MockRM(conf);
    rm.init(conf);
    stateStore = rm.getFederationStateStoreService().getStateStoreClient();

    // generate an [app01] and join the [SC-1] cluster.
    List<ApplicationId> appIds = new ArrayList<>();
    ApplicationId appId01 = ApplicationId.newInstance(Time.now(), 1);
    addApplication2StateStore(appId01, stateStore);
    appIds.add(appId01);

    // generate an [app02] and join the [SC-1] cluster.
    ApplicationId appId02 = ApplicationId.newInstance(Time.now(), 2);
    addApplication2StateStore(appId02, stateStore);
    appIds.add(appId02);

    // generate an [app03] and join the [SC-1] cluster.
    ApplicationId appId03 = ApplicationId.newInstance(Time.now(), 3);
    addApplication2StateStore(appId03, stateStore);
    appIds.add(appId03);

    // make sure the apps can be queried in the stateStore.
    GetApplicationsHomeSubClusterRequest allRequest =
        GetApplicationsHomeSubClusterRequest.newInstance(subClusterId);
    GetApplicationsHomeSubClusterResponse allResponse =
        stateStore.getApplicationsHomeSubCluster(allRequest);
    Assert.assertNotNull(allResponse);
    List<ApplicationHomeSubCluster> appHomeSCLists = allResponse.getAppsHomeSubClusters();
    Assert.assertNotNull(appHomeSCLists);
    Assert.assertEquals(3, appHomeSCLists.size());

    // app04 exists in both RM memory and stateStore.
    ApplicationId appId04 = ApplicationId.newInstance(Time.now(), 4);
    addApplication2StateStore(appId04, stateStore);
    addApplication2RMAppManager(rm, appId04);

    // start rm.
    rm.start();

    // wait 5s, wait for the thread to finish cleaning up.
    GenericTestUtils.waitFor(() -> {
      int appsSize = 0;
      try {
        List<ApplicationHomeSubCluster> subClusters =
            getApplicationsFromStateStore();
        Assert.assertNotNull(subClusters);
        appsSize = subClusters.size();
      } catch (YarnException e) {
        e.printStackTrace();
      }
      return (appsSize == 1);
    }, 100, 1000 * 5);

    // check the app to make sure the apps(app01,app02,app03) doesn't exist.
    for (ApplicationId appId : appIds) {
      GetApplicationHomeSubClusterRequest appRequest =
          GetApplicationHomeSubClusterRequest.newInstance(appId);
      LambdaTestUtils.intercept(FederationStateStoreException.class,
          "Application " + appId + " does not exist",
          () -> stateStore.getApplicationHomeSubCluster(appRequest));
    }

    if (rm != null) {
      rm.stop();
      rm = null;
    }
  }

  @Test
  public void testCleanUpApplicationWhenRMCompleteOneApp() throws Exception {

    // We design such a test case.
    // Step1. We start RM，Set the RM memory to keep a maximum of 1 completed app.
    // Step2. Register app[01-03] to RM memory & stateStore.
    // Step3. We clean up app01, app02, app03, at this time,
    // app01, app02 should be cleaned up from statestore, app03 should remain in statestore.

    // set yarn configuration.
    conf.setBoolean(YarnConfiguration.FEDERATION_ENABLED, true);
    conf.setInt(YarnConfiguration.FEDERATION_STATESTORE_HEARTBEAT_INITIAL_DELAY, 10);
    conf.set(YarnConfiguration.RM_CLUSTER_ID, subClusterId.getId());
    conf.setBoolean(YarnConfiguration.RECOVERY_ENABLED, true);
    conf.setInt(YarnConfiguration.RM_STATE_STORE_MAX_COMPLETED_APPLICATIONS, 1);
    conf.set(YarnConfiguration.RM_STORE, MemoryRMStateStore.class.getName());

    // set up MockRM.
    MockRM rm = new MockRM(conf);
    rm.init(conf);
    stateStore = rm.getFederationStateStoreService().getStateStoreClient();
    rm.start();

    // generate an [app01] and join the [SC-1] cluster.
    List<ApplicationId> appIds = new ArrayList<>();
    ApplicationId appId01 = ApplicationId.newInstance(Time.now(), 1);
    addApplication2StateStore(appId01, stateStore);
    addApplication2RMAppManager(rm, appId01);
    appIds.add(appId01);

    // generate an [app02] and join the [SC-1] cluster.
    ApplicationId appId02 = ApplicationId.newInstance(Time.now(), 2);
    addApplication2StateStore(appId02, stateStore);
    addApplication2RMAppManager(rm, appId02);
    appIds.add(appId02);

    // generate an [app03] and join the [SC-1] cluster.
    ApplicationId appId03 = ApplicationId.newInstance(Time.now(), 3);
    addApplication2StateStore(appId03, stateStore);
    addApplication2RMAppManager(rm, appId03);

    // rmAppManager
    RMAppManager rmAppManager = rm.getRMAppManager();
    rmAppManager.finishApplication4Test(appId01);
    rmAppManager.finishApplication4Test(appId02);
    rmAppManager.finishApplication4Test(appId03);
    rmAppManager.checkAppNumCompletedLimit4Test();

    // app01, app02 should be cleaned from statestore
    // After the query, it should report the error not exist.
    for (ApplicationId appId : appIds) {
      GetApplicationHomeSubClusterRequest appRequest =
          GetApplicationHomeSubClusterRequest.newInstance(appId);
      LambdaTestUtils.intercept(FederationStateStoreException.class,
          "Application " + appId + " does not exist",
          () -> stateStore.getApplicationHomeSubCluster(appRequest));
    }

    // app03 should remain in statestore
    List<ApplicationHomeSubCluster> appHomeScList = getApplicationsFromStateStore();
    Assert.assertNotNull(appHomeScList);
    Assert.assertEquals(1, appHomeScList.size());
    ApplicationHomeSubCluster homeSubCluster = appHomeScList.get(0);
    Assert.assertNotNull(homeSubCluster);
    Assert.assertEquals(appId03, homeSubCluster.getApplicationId());
  }

  private void addApplication2StateStore(ApplicationId appId,
      FederationStateStore fedStateStore) throws YarnException {
    ApplicationHomeSubCluster appHomeSC = ApplicationHomeSubCluster.newInstance(
        appId, subClusterId);
    AddApplicationHomeSubClusterRequest addHomeSCRequest =
        AddApplicationHomeSubClusterRequest.newInstance(appHomeSC);
    fedStateStore.addApplicationHomeSubCluster(addHomeSCRequest);
  }

  private List<ApplicationHomeSubCluster> getApplicationsFromStateStore() throws YarnException {
    // make sure the apps can be queried in the stateStore
    GetApplicationsHomeSubClusterRequest allRequest =
        GetApplicationsHomeSubClusterRequest.newInstance(subClusterId);
    GetApplicationsHomeSubClusterResponse allResponse =
        stateStore.getApplicationsHomeSubCluster(allRequest);
    Assert.assertNotNull(allResponse);
    List<ApplicationHomeSubCluster> appHomeSCLists = allResponse.getAppsHomeSubClusters();
    Assert.assertNotNull(appHomeSCLists);
    return appHomeSCLists;
  }

  private void addApplication2RMAppManager(MockRM rm, ApplicationId appId) {
    RMContext rmContext = rm.getRMContext();
    Map<ApplicationId, RMApp> rmAppMaps = rmContext.getRMApps();
    String user = MockApps.newUserName();
    String name = MockApps.newAppName();
    String queue = MockApps.newQueue();

    YarnScheduler scheduler = mock(YarnScheduler.class);

    ApplicationMasterService masterService =
        new ApplicationMasterService(rmContext, scheduler);

    ApplicationSubmissionContext submissionContext =
        new ApplicationSubmissionContextPBImpl();

    // applicationId will not be used because RMStateStore is mocked,
    // but applicationId is still set for safety
    submissionContext.setApplicationId(appId);
    submissionContext.setPriority(Priority.newInstance(0));

    RMApp application = new RMAppImpl(appId, rmContext, conf, name,
        user, queue, submissionContext, scheduler, masterService,
        System.currentTimeMillis(), "YARN", null,
        new ArrayList<>());

    rmAppMaps.putIfAbsent(application.getApplicationId(), application);
  }


  @Test
  public void testPolicyConfigurationMethod() throws YarnException {

    // This test case tests 3 methods.
    // 1.setPolicyConfiguration
    // 2.getPolicyConfiguration
    // 3.getPolicyConfigurations
    FederationStateStoreService stateStoreService = mockRM.getFederationStateStoreService();

    // set queue basic information (queue1)
    String queue1 = "queue1";
    SubClusterPolicyConfiguration requestPolicyConf1 = getUniformPolicy(queue1);
    SetSubClusterPolicyConfigurationRequest configurationRequest1 =
        SetSubClusterPolicyConfigurationRequest.newInstance(requestPolicyConf1);
    // store policy configuration (queue1)
    stateStoreService.setPolicyConfiguration(configurationRequest1);

    // set queue basic information (queue2)
    String queue2 = "queue2";
    SubClusterPolicyConfiguration requestPolicyConf2 = getUniformPolicy(queue2);
    SetSubClusterPolicyConfigurationRequest configurationRequest2 =
        SetSubClusterPolicyConfigurationRequest.newInstance(requestPolicyConf2);
    // store policy configuration (queue1)
    stateStoreService.setPolicyConfiguration(configurationRequest2);

    // get policy configuration
    GetSubClusterPolicyConfigurationRequest request1 =
        GetSubClusterPolicyConfigurationRequest.newInstance(queue1);
    GetSubClusterPolicyConfigurationResponse response =
        stateStoreService.getPolicyConfiguration(request1);
    Assert.assertNotNull(response);

    SubClusterPolicyConfiguration responsePolicyConf =
        response.getPolicyConfiguration();
    Assert.assertNotNull(responsePolicyConf);
    Assert.assertEquals(requestPolicyConf1, responsePolicyConf);

    // get policy configurations
    GetSubClusterPoliciesConfigurationsRequest policiesRequest1 =
        GetSubClusterPoliciesConfigurationsRequest.newInstance();
    GetSubClusterPoliciesConfigurationsResponse policiesResponse1 =
        stateStoreService.getPoliciesConfigurations(policiesRequest1);
    Assert.assertNotNull(policiesResponse1);

    List<SubClusterPolicyConfiguration> policiesConfigs = policiesResponse1.getPoliciesConfigs();
    Assert.assertNotNull(policiesConfigs);
    Assert.assertEquals(2, policiesConfigs.size());
    Assert.assertTrue(policiesConfigs.contains(requestPolicyConf1));
    Assert.assertTrue(policiesConfigs.contains(requestPolicyConf2));
  }

  public SubClusterPolicyConfiguration getUniformPolicy(String queue)
      throws FederationPolicyInitializationException {
    UniformBroadcastPolicyManager wfp = new UniformBroadcastPolicyManager();
    wfp.setQueue(queue);
    SubClusterPolicyConfiguration fpc = wfp.serializeConf();
    return fpc;
  }

  @Test
  public void testSubClusterMethod() throws YarnException {

    // This test case tests 5 methods.
    // 1.registerSubCluster
    // 2.deregisterSubCluster
    // 3.subClusterHeartbeat
    // 4.getSubCluster
    // 5.getSubClusters

    FederationStateStoreService stateStoreService =
        mockRM.getFederationStateStoreService();

    // registerSubCluster subCluster1
    SubClusterId subClusterId1 = SubClusterId.newInstance("SC1");
    SubClusterInfo subClusterInfo1 = createSubClusterInfo(subClusterId1);

    SubClusterRegisterRequest registerRequest1 =
        SubClusterRegisterRequest.newInstance(subClusterInfo1);
    stateStoreService.registerSubCluster(registerRequest1);

    // registerSubCluster subCluster2
    SubClusterId subClusterId2 = SubClusterId.newInstance("SC2");
    SubClusterInfo subClusterInfo2 = createSubClusterInfo(subClusterId2);

    SubClusterRegisterRequest registerRequest2 =
        SubClusterRegisterRequest.newInstance(subClusterInfo2);
    stateStoreService.registerSubCluster(registerRequest2);

    // getSubCluster subCluster1
    GetSubClusterInfoRequest subClusterRequest =
        GetSubClusterInfoRequest.newInstance(subClusterId1);
    GetSubClusterInfoResponse subClusterResponse =
        stateStoreService.getSubCluster(subClusterRequest);
    Assert.assertNotNull(subClusterResponse);

    // We query subCluster1, we want to get SubClusterInfo of subCluster1
    SubClusterInfo subClusterInfo1Resp = subClusterResponse.getSubClusterInfo();
    Assert.assertNotNull(subClusterInfo1Resp);
    Assert.assertEquals(subClusterInfo1, subClusterInfo1Resp);

    // We call the getSubClusters method and filter the Active SubCluster
    // subCluster1 and subCluster2 are just registered, they are in NEW state,
    // so we will get 0 active subclusters
    GetSubClustersInfoRequest subClustersInfoRequest =
        GetSubClustersInfoRequest.newInstance(true);
    GetSubClustersInfoResponse subClustersInfoResp =
        stateStoreService.getSubClusters(subClustersInfoRequest);
    Assert.assertNotNull(subClustersInfoResp);
    List<SubClusterInfo> subClusterInfos = subClustersInfoResp.getSubClusters();
    Assert.assertNotNull(subClusterInfos);
    Assert.assertEquals(0, subClusterInfos.size());

    // We let subCluster1 heartbeat and set subCluster1 to Running state
    SubClusterHeartbeatRequest heartbeatRequest =
        SubClusterHeartbeatRequest.newInstance(subClusterId1, SubClusterState.SC_RUNNING,
        "capability");
    SubClusterHeartbeatResponse heartbeatResponse =
        stateStoreService.subClusterHeartbeat(heartbeatRequest);
    Assert.assertNotNull(heartbeatResponse);

    // We call the getSubClusters method again and filter the Active SubCluster
    // We want to get 1 active SubCluster
    GetSubClustersInfoRequest subClustersInfoRequest1 =
        GetSubClustersInfoRequest.newInstance(true);
    GetSubClustersInfoResponse subClustersInfoResp1 =
        stateStoreService.getSubClusters(subClustersInfoRequest1);
    Assert.assertNotNull(subClustersInfoResp1);
    List<SubClusterInfo> subClusterInfos1 = subClustersInfoResp1.getSubClusters();
    Assert.assertNotNull(subClusterInfos1);
    Assert.assertEquals(1, subClusterInfos1.size());
  }

  private SubClusterInfo createSubClusterInfo(SubClusterId clusterId) {

    String amRMAddress = "1.2.3.4:1";
    String clientRMAddress = "1.2.3.4:2";
    String rmAdminAddress = "1.2.3.4:3";
    String webAppAddress = "1.2.3.4:4";

    return SubClusterInfo.newInstance(clusterId, amRMAddress,
        clientRMAddress, rmAdminAddress, webAppAddress, SubClusterState.SC_NEW,
        Time.now(), "capability");
  }
}