http://git-wip-us.apache.org/repos/asf/hadoop/blob/28c8a586/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/test/java/org/apache/hadoop/yarn/server/federation/utils/FederationPoliciesTestUtil.java ---------------------------------------------------------------------- diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/test/java/org/apache/hadoop/yarn/server/federation/utils/FederationPoliciesTestUtil.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/test/java/org/apache/hadoop/yarn/server/federation/utils/FederationPoliciesTestUtil.java index 8c2115b..f901329 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/test/java/org/apache/hadoop/yarn/server/federation/utils/FederationPoliciesTestUtil.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/test/java/org/apache/hadoop/yarn/server/federation/utils/FederationPoliciesTestUtil.java @@ -19,13 +19,20 @@ package org.apache.hadoop.yarn.server.federation.utils; import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.yarn.conf.YarnConfiguration; import org.apache.hadoop.yarn.exceptions.YarnException; +import org.apache.hadoop.yarn.server.federation.policies.ConfigurableFederationPolicy; +import org.apache.hadoop.yarn.server.federation.policies.FederationPolicyInitializationContext; +import org.apache.hadoop.yarn.server.federation.policies.dao.WeightedPolicyInfo; import org.apache.hadoop.yarn.server.federation.resolver.DefaultSubClusterResolverImpl; import org.apache.hadoop.yarn.server.federation.resolver.SubClusterResolver; import org.apache.hadoop.yarn.server.federation.store.FederationStateStore; -import org.apache.hadoop.yarn.server.federation.store.records.GetSubClustersInfoResponse; +import org.apache.hadoop.yarn.server.federation.store.records.*; import java.net.URL; +import java.nio.ByteBuffer; import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; import static org.mockito.Matchers.any; import static org.mockito.Mockito.mock; @@ -41,6 +48,41 @@ public final class FederationPoliciesTestUtil { // disabled. } + + public static void initializePolicyContext( + FederationPolicyInitializationContext fpc, ConfigurableFederationPolicy + policy, WeightedPolicyInfo policyInfo, + Map<SubClusterId, SubClusterInfo> activeSubclusters) + throws YarnException { + ByteBuffer buf = policyInfo.toByteBuffer(); + fpc.setSubClusterPolicyConfiguration(SubClusterPolicyConfiguration + .newInstance("queue1", policy.getClass().getCanonicalName(), buf)); + FederationStateStoreFacade facade = FederationStateStoreFacade + .getInstance(); + FederationStateStore fss = mock(FederationStateStore.class); + + if (activeSubclusters == null) { + activeSubclusters = new HashMap<SubClusterId, SubClusterInfo>(); + } + GetSubClustersInfoResponse response = GetSubClustersInfoResponse + .newInstance(new ArrayList<SubClusterInfo>(activeSubclusters.values())); + + when(fss.getSubClusters(any())).thenReturn(response); + facade.reinitialize(fss, new Configuration()); + fpc.setFederationStateStoreFacade(facade); + policy.reinitialize(fpc); + } + + public static void initializePolicyContext( + ConfigurableFederationPolicy policy, + WeightedPolicyInfo policyInfo, Map<SubClusterId, + SubClusterInfo> activeSubclusters) throws YarnException { + FederationPolicyInitializationContext context = + new FederationPolicyInitializationContext(null, initResolver(), + initFacade()); + initializePolicyContext(context, policy, policyInfo, activeSubclusters); + } + /** * Initialize a {@link SubClusterResolver}. * @@ -66,18 +108,52 @@ public final class FederationPoliciesTestUtil { * Initialiaze a main-memory {@link FederationStateStoreFacade} used for * testing, wiht a mock resolver. * + * @param subClusterInfos the list of subclusters to be served on + * getSubClusters invocations. + * * @return the facade. * * @throws YarnException in case the initialization is not successful. */ - public static FederationStateStoreFacade initFacade() throws YarnException { + + public static FederationStateStoreFacade initFacade( + List<SubClusterInfo> subClusterInfos, SubClusterPolicyConfiguration + policyConfiguration) throws YarnException { FederationStateStoreFacade goodFacade = FederationStateStoreFacade .getInstance(); FederationStateStore fss = mock(FederationStateStore.class); GetSubClustersInfoResponse response = GetSubClustersInfoResponse - .newInstance(new ArrayList<>()); + .newInstance(subClusterInfos); when(fss.getSubClusters(any())).thenReturn(response); + + List<SubClusterPolicyConfiguration> configurations = new ArrayList<>(); + configurations.add(policyConfiguration); + + GetSubClusterPoliciesConfigurationsResponse policiesResponse = + GetSubClusterPoliciesConfigurationsResponse + .newInstance(configurations); + when(fss.getPoliciesConfigurations(any())).thenReturn(policiesResponse); + + GetSubClusterPolicyConfigurationResponse policyResponse = + GetSubClusterPolicyConfigurationResponse + .newInstance(policyConfiguration); + when(fss.getPolicyConfiguration(any())).thenReturn(policyResponse); + goodFacade.reinitialize(fss, new Configuration()); return goodFacade; } + + /** + * Initialiaze a main-memory {@link FederationStateStoreFacade} used for + * testing, wiht a mock resolver. + * + * @return the facade. + * + * @throws YarnException in case the initialization is not successful. + */ + public static FederationStateStoreFacade initFacade() throws YarnException { + return initFacade(new ArrayList<>(), mock(SubClusterPolicyConfiguration + .class)); + } + }
--------------------------------------------------------------------- To unsubscribe, e-mail: common-commits-unsubscr...@hadoop.apache.org For additional commands, e-mail: common-commits-h...@hadoop.apache.org