YARN-11324. Fix some PBImpl classes to avoid NPE. (#4961)

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/federation/store/impl/MemoryFederationStateStore.java

@@ -33,7 +33,6 @@
 import org.apache.hadoop.classification.VisibleForTesting;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.security.token.delegation.DelegationKey;
-import org.apache.hadoop.util.Time;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
 import org.apache.hadoop.yarn.api.records.ReservationId;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
@@ -295,7 +294,7 @@ public GetApplicationsHomeSubClusterResponse getApplicationsHomeSubCluster(
 
   private ApplicationHomeSubCluster generateAppHomeSC(ApplicationId applicationId) {
     SubClusterId subClusterId = applications.get(applicationId);
-    return ApplicationHomeSubCluster.newInstance(applicationId, Time.now(), subClusterId);
+    return ApplicationHomeSubCluster.newInstance(applicationId, subClusterId);
   }
 
   @Override

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/federation/store/impl/ZookeeperFederationStateStore.java

@@ -30,7 +30,6 @@
 import org.apache.commons.lang3.NotImplementedException;
 import org.apache.hadoop.classification.VisibleForTesting;
 import org.apache.hadoop.conf.Configuration;
-import org.apache.hadoop.util.Time;
 import org.apache.hadoop.util.curator.ZKCuratorManager;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
@@ -300,7 +299,7 @@ private ApplicationHomeSubCluster generateAppHomeSC(String appId) {
       ApplicationId applicationId = ApplicationId.fromString(appId);
       SubClusterId homeSubCluster = getApp(applicationId);
       ApplicationHomeSubCluster app =
-          ApplicationHomeSubCluster.newInstance(applicationId, Time.now(), homeSubCluster);
+          ApplicationHomeSubCluster.newInstance(applicationId, homeSubCluster);
       return app;
     } catch (Exception ex) {
       LOG.error("get homeSubCluster by appId = {}.", appId);

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/federation/store/records/impl/pb/ApplicationHomeSubClusterPBImpl.java

@@ -44,6 +44,7 @@ public class ApplicationHomeSubClusterPBImpl extends ApplicationHomeSubCluster {
 
   private ApplicationId applicationId = null;
   private SubClusterId homeSubCluster = null;
+  private long createTime = 0L;
 
   public ApplicationHomeSubClusterPBImpl() {
     builder = ApplicationHomeSubClusterProto.newBuilder();
@@ -110,6 +111,9 @@ public String toString() {
   @Override
   public ApplicationId getApplicationId() {
     ApplicationHomeSubClusterProtoOrBuilder p = viaProto ? proto : builder;
+    if (this.applicationId != null) {
+      return this.applicationId;
+    }
     if (!p.hasApplicationId()) {
       return null;
     }
@@ -125,6 +129,7 @@ public void setApplicationId(ApplicationId applicationId) {
       return;
     }
     this.applicationId = applicationId;
+    builder.setApplicationId(convertToProtoFormat(applicationId));
   }
 
   @Override
@@ -141,22 +146,34 @@ public SubClusterId getHomeSubCluster() {
   }
 
   @Override
-  public void setHomeSubCluster(SubClusterId homeSubCluster) {
+  public void setHomeSubCluster(SubClusterId paramHomeSubCluster) {
     maybeInitBuilder();
-    if (homeSubCluster == null) {
+    if (paramHomeSubCluster == null) {
       builder.clearHomeSubCluster();
+      return;
     }
-    this.homeSubCluster = homeSubCluster;
+    this.homeSubCluster = paramHomeSubCluster;
+    builder.setHomeSubCluster(convertToProtoFormat(paramHomeSubCluster));
   }
 
   @Override
   public long getCreateTime() {
-    return 0;
+    ApplicationHomeSubClusterProtoOrBuilder p = viaProto ? proto : builder;
+    if (this.createTime != 0) {
+      return this.createTime;
+    }
+    if (!p.hasCreateTime()) {
+      return 0;
+    }
+    this.createTime = p.getCreateTime();
+    return this.createTime;
   }
 
   @Override
   public void setCreateTime(long time) {
-
+    maybeInitBuilder();
+    this.createTime = time;
+    builder.setCreateTime(time);
   }
 
   private SubClusterId convertFromProtoFormat(SubClusterIdProto subClusterId) {

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/federation/store/records/impl/pb/GetApplicationHomeSubClusterRequestPBImpl.java

@@ -127,6 +127,7 @@ public void setApplicationId(ApplicationId applicationId) {
       return;
     }
     this.applicationId = applicationId;
+    builder.setApplicationId(convertToProtoFormat(applicationId));
   }
 
   private ApplicationId convertFromProtoFormat(ApplicationIdProto appId) {

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/federation/store/records/impl/pb/GetApplicationsHomeSubClusterResponsePBImpl.java

@@ -122,6 +122,7 @@ public void setAppsHomeSubClusters(
       return;
     }
     this.appsHomeSubCluster = appsHomeSubClusters;
+    addSubClustersInfoToProto();
   }
 
   private void initSubClustersInfoList() {
@@ -132,7 +133,7 @@ private void initSubClustersInfoList() {
         viaProto ? proto : builder;
     List<ApplicationHomeSubClusterProto> subClusterInfosList =
         p.getAppSubclusterMapList();
-    appsHomeSubCluster = new ArrayList<ApplicationHomeSubCluster>();
+    appsHomeSubCluster = new ArrayList<>();
 
     for (ApplicationHomeSubClusterProto r : subClusterInfosList) {
       appsHomeSubCluster.add(convertFromProtoFormat(r));

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/federation/store/records/impl/pb/GetReservationHomeSubClusterRequestPBImpl.java

@@ -119,13 +119,14 @@ public ReservationId getReservationId() {
   }
 
   @Override
-  public void setReservationId(ReservationId reservationId) {
+  public void setReservationId(ReservationId paramReservationId) {
     maybeInitBuilder();
-    if (reservationId == null) {
+    if (paramReservationId == null) {
       builder.clearReservationId();
       return;
     }
-    this.reservationId = reservationId;
+    this.reservationId = paramReservationId;
+    builder.setReservationId(convertToProtoFormat(paramReservationId));
   }
 
   private ReservationId convertFromProtoFormat(ReservationIdProto appId) {

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/federation/store/records/impl/pb/GetReservationsHomeSubClusterResponsePBImpl.java

@@ -122,6 +122,7 @@ public void setAppsHomeSubClusters(
       return;
     }
     this.appsHomeSubCluster = appsHomeSubClusters;
+    addSubClustersInfoToProto();
   }
 
   private void initSubClustersInfoList() {

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/federation/store/records/impl/pb/GetSubClusterInfoResponsePBImpl.java

@@ -114,12 +114,14 @@ public SubClusterInfo getSubClusterInfo() {
   }
 
   @Override
-  public void setSubClusterInfo(SubClusterInfo subClusterInfo) {
+  public void setSubClusterInfo(SubClusterInfo paramSubClusterInfo) {
     maybeInitBuilder();
-    if (subClusterInfo == null) {
+    if (paramSubClusterInfo == null) {
       builder.clearSubClusterInfo();
+      return;
     }
-    this.subClusterInfo = subClusterInfo;
+    this.subClusterInfo = paramSubClusterInfo;
+    builder.setSubClusterInfo(convertToProtoFormat(paramSubClusterInfo));
   }
 
   private SubClusterInfo convertFromProtoFormat(

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/federation/store/records/impl/pb/GetSubClusterPoliciesConfigurationsResponsePBImpl.java

@@ -123,6 +123,7 @@ public void setPoliciesConfigs(
       builder.clearPoliciesConfigurations();
     }
     this.subClusterPolicies = policyConfigurations;
+    addSubClusterPoliciesConfigurationsToProto();
   }
 
   private void initSubClusterPoliciesConfigurationsList() {

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/federation/store/records/impl/pb/GetSubClusterPolicyConfigurationResponsePBImpl.java

@@ -126,8 +126,10 @@ public void setPolicyConfiguration(
     maybeInitBuilder();
     if (policyConfiguration == null) {
       builder.clearPolicyConfiguration();
+      return;
     }
     this.subClusterPolicy = policyConfiguration;
+    this.builder.setPolicyConfiguration(convertToProtoFormat(policyConfiguration));
   }
 
   private SubClusterPolicyConfiguration convertFromProtoFormat(

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/federation/store/records/impl/pb/GetSubClustersInfoResponsePBImpl.java

@@ -101,6 +101,7 @@ public void setSubClusters(Collection<SubClusterInfo> subClusters) {
       return;
     }
     this.subClusterInfos = subClusters.stream().collect(Collectors.toList());
+    addSubClusterInfosToProto();
   }
 
   private void initSubClustersInfoList() {

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/federation/store/records/impl/pb/ReservationHomeSubClusterPBImpl.java

@@ -124,6 +124,7 @@ public void setReservationId(ReservationId resId) {
       builder.clearReservationId();
       return;
     }
+    builder.setReservationId(convertToProtoFormat(resId));
     this.reservationId = resId;
   }
 

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/federation/store/records/impl/pb/RouterMasterKeyRequestPBImpl.java

@@ -93,8 +93,10 @@ public void setRouterMasterKey(RouterMasterKey masterKey) {
     maybeInitBuilder();
     if (masterKey == null) {
       builder.clearRouterMasterKey();
+      return;
     }
     this.routerMasterKey = masterKey;
+    builder.setRouterMasterKey(convertToProtoFormat(masterKey));
   }
 
   @Override

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/federation/store/records/impl/pb/RouterMasterKeyResponsePBImpl.java

@@ -93,8 +93,10 @@ public void setRouterMasterKey(RouterMasterKey masterKey) {
     maybeInitBuilder();
     if (masterKey == null) {
       builder.clearRouterMasterKey();
+      return;
     }
     this.routerMasterKey = masterKey;
+    builder.setRouterMasterKey(convertToProtoFormat(masterKey));
   }
 
   @Override

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/federation/store/records/impl/pb/SetSubClusterPolicyConfigurationRequestPBImpl.java

@@ -126,8 +126,10 @@ public void setPolicyConfiguration(
     maybeInitBuilder();
     if (policyConfiguration == null) {
       builder.clearPolicyConfiguration();
+      return;
     }
     this.subClusterPolicy = policyConfiguration;
+    builder.setPolicyConfiguration(convertToProtoFormat(policyConfiguration));
   }
 
   private SubClusterPolicyConfiguration convertFromProtoFormat(

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/federation/store/records/impl/pb/SubClusterHeartbeatRequestPBImpl.java

@@ -122,8 +122,10 @@ public void setSubClusterId(SubClusterId subClusterId) {
     maybeInitBuilder();
     if (subClusterId == null) {
       builder.clearSubClusterId();
+      return;
     }
     this.subClusterId = subClusterId;
+    builder.setSubClusterId(convertToProtoFormat(subClusterId));
   }
 
   @Override

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/federation/store/records/impl/pb/SubClusterInfoPBImpl.java

@@ -104,8 +104,10 @@ public void setSubClusterId(SubClusterId subClusterId) {
     maybeInitBuilder();
     if (subClusterId == null) {
       builder.clearSubClusterId();
+      return;
     }
     this.subClusterId = subClusterId;
+    builder.setSubClusterId(convertToProtoFormat(subClusterId));
   }
 
   @Override

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/federation/store/records/impl/pb/SubClusterRegisterRequestPBImpl.java

@@ -118,8 +118,10 @@ public void setSubClusterInfo(SubClusterInfo subClusterInfo) {
     maybeInitBuilder();
     if (subClusterInfo == null) {
       builder.clearSubClusterInfo();
+      return;
     }
     this.subClusterInfo = subClusterInfo;
+    builder.setSubClusterInfo(convertToProtoFormat(subClusterInfo));
   }
 
   private SubClusterInfo convertFromProtoFormat(

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/test/java/org/apache/hadoop/yarn/server/federation/policies/amrmproxy/TestBroadcastAMRMProxyFederationPolicy.java

@@ -19,7 +19,6 @@
 package org.apache.hadoop.yarn.server.federation.policies.amrmproxy;
 
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
 
 import java.util.HashSet;
 import java.util.List;
@@ -52,9 +51,9 @@ public void setUp() throws Exception {
 
     for (int i = 1; i <= 2; i++) {
       SubClusterIdInfo sc = new SubClusterIdInfo("sc" + i);
-      SubClusterInfo sci = mock(SubClusterInfo.class);
+      SubClusterInfo sci = SubClusterInfo.newInstance(
-      when(sci.getState()).thenReturn(SubClusterState.SC_RUNNING);
+          sc.toId(), "dns1:80", "dns1:81", "dns1:82", "dns1:83", SubClusterState.SC_RUNNING,
-      when(sci.getSubClusterId()).thenReturn(sc.toId());
+          System.currentTimeMillis(), "something");
       getActiveSubclusters().put(sc.toId(), sci);
     }
 

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/test/java/org/apache/hadoop/yarn/server/federation/policies/amrmproxy/TestHomeAMRMProxyPolicy.java

@@ -24,7 +24,6 @@
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.fail;
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
 
 import java.util.HashSet;
 import java.util.List;
@@ -62,9 +61,9 @@ public void setUp() throws Exception {
 
     for (int i = 0; i < NUM_SUBCLUSTERS; i++) {
       SubClusterIdInfo sc = new SubClusterIdInfo("sc" + i);
-      SubClusterInfo sci = mock(SubClusterInfo.class);
+      SubClusterInfo sci = SubClusterInfo.newInstance(
-      when(sci.getState()).thenReturn(SubClusterState.SC_RUNNING);
+          sc.toId(), "dns1:80", "dns1:81", "dns1:82", "dns1:83", SubClusterState.SC_RUNNING,
-      when(sci.getSubClusterId()).thenReturn(sc.toId());
+          System.currentTimeMillis(), "something");
       getActiveSubclusters().put(sc.toId(), sci);
     }
 

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/test/java/org/apache/hadoop/yarn/server/federation/policies/amrmproxy/TestLocalityMulticastAMRMProxyPolicy.java

@@ -19,8 +19,6 @@
 package org.apache.hadoop.yarn.server.federation.policies.amrmproxy;
 
 import static org.junit.Assert.fail;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
 
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
@@ -79,9 +77,9 @@ public void setUp() throws Exception {
       SubClusterIdInfo sc = new SubClusterIdInfo("subcluster" + i);
       // sub-cluster 3 is not active
       if (i != 3) {
-        SubClusterInfo sci = mock(SubClusterInfo.class);
+        SubClusterInfo sci = SubClusterInfo.newInstance(
-        when(sci.getState()).thenReturn(SubClusterState.SC_RUNNING);
+            sc.toId(), "dns1:80", "dns1:81", "dns1:82", "dns1:83", SubClusterState.SC_RUNNING,
-        when(sci.getSubClusterId()).thenReturn(sc.toId());
+            System.currentTimeMillis(), "something");
         getActiveSubclusters().put(sc.toId(), sci);
       }
 
@@ -330,11 +328,14 @@ Collections.<NodeReport> emptyList(),
    * use as the default for when nodes or racks are unknown.
    */
   private void addHomeSubClusterAsActive() {
-    SubClusterInfo sci = mock(SubClusterInfo.class);
+
-    when(sci.getState()).thenReturn(SubClusterState.SC_RUNNING);
+    SubClusterId homeSubCluster = getHomeSubCluster();
-    when(sci.getSubClusterId()).thenReturn(getHomeSubCluster());
+    SubClusterInfo sci = SubClusterInfo.newInstance(
-    getActiveSubclusters().put(getHomeSubCluster(), sci);
+        homeSubCluster, "dns1:80", "dns1:81", "dns1:82", "dns1:83", SubClusterState.SC_RUNNING,
-    SubClusterIdInfo sc = new SubClusterIdInfo(getHomeSubCluster().getId());
+        System.currentTimeMillis(), "something");
+
+    getActiveSubclusters().put(homeSubCluster, sci);
+    SubClusterIdInfo sc = new SubClusterIdInfo(homeSubCluster.getId());
 
     getPolicyInfo().getRouterPolicyWeights().put(sc, 0.1f);
     getPolicyInfo().getAMRMPolicyWeights().put(sc, 0.1f);

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/test/java/org/apache/hadoop/yarn/server/federation/policies/amrmproxy/TestRejectAMRMProxyPolicy.java

@@ -19,7 +19,6 @@
 package org.apache.hadoop.yarn.server.federation.policies.amrmproxy;
 
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
 
 import java.util.HashSet;
 import java.util.List;
@@ -51,9 +50,9 @@ public void setUp() throws Exception {
 
     for (int i = 1; i <= 2; i++) {
       SubClusterIdInfo sc = new SubClusterIdInfo("sc" + i);
-      SubClusterInfo sci = mock(SubClusterInfo.class);
+      SubClusterInfo sci = SubClusterInfo.newInstance(
-      when(sci.getState()).thenReturn(SubClusterState.SC_RUNNING);
+          sc.toId(), "dns1:80", "dns1:81", "dns1:82", "dns1:83", SubClusterState.SC_RUNNING,
-      when(sci.getSubClusterId()).thenReturn(sc.toId());
+          System.currentTimeMillis(), "something");
       getActiveSubclusters().put(sc.toId(), sci);
     }
 

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/test/java/org/apache/hadoop/yarn/server/federation/policies/router/TestPriorityRouterPolicy.java

@@ -17,8 +17,6 @@
 package org.apache.hadoop.yarn.server.federation.policies.router;
 
 import static org.apache.hadoop.test.LambdaTestUtils.intercept;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
 
 import java.util.HashMap;
 import java.util.Map;
@@ -93,9 +91,9 @@ public void testZeroSubClustersWithPositiveWeight() throws Exception {
     for (int i = 0; i < 5; i++) {
       SubClusterIdInfo sc = new SubClusterIdInfo("sc" + i);
 
-      SubClusterInfo sci = mock(SubClusterInfo.class);
+      SubClusterInfo sci = SubClusterInfo.newInstance(
-      when(sci.getState()).thenReturn(SubClusterState.SC_RUNNING);
+          sc.toId(), "dns1:80", "dns1:81", "dns1:82", "dns1:83", SubClusterState.SC_RUNNING,
-      when(sci.getSubClusterId()).thenReturn(sc.toId());
+          System.currentTimeMillis(), "something");
       getActiveSubclusters().put(sc.toId(), sci);
       routerWeights.put(sc, 0.0f);
       amrmWeights.put(sc, -1.0f);

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/test/java/org/apache/hadoop/yarn/server/federation/store/records/TestFederationProtocolRecords.java

@@ -19,6 +19,7 @@
 
 import org.apache.hadoop.yarn.api.BasePBImplRecordsTest;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
+import org.apache.hadoop.yarn.api.records.ReservationId;
 import org.apache.hadoop.yarn.federation.proto.YarnServerFederationProtos.AddApplicationHomeSubClusterRequestProto;
 import org.apache.hadoop.yarn.federation.proto.YarnServerFederationProtos.AddApplicationHomeSubClusterResponseProto;
 import org.apache.hadoop.yarn.federation.proto.YarnServerFederationProtos.DeleteApplicationHomeSubClusterRequestProto;
@@ -47,9 +48,11 @@
 import org.apache.hadoop.yarn.federation.proto.YarnServerFederationProtos.SubClusterRegisterResponseProto;
 import org.apache.hadoop.yarn.federation.proto.YarnServerFederationProtos.UpdateApplicationHomeSubClusterRequestProto;
 import org.apache.hadoop.yarn.federation.proto.YarnServerFederationProtos.UpdateApplicationHomeSubClusterResponseProto;
+import org.apache.hadoop.yarn.federation.proto.YarnServerFederationProtos.GetReservationHomeSubClusterRequestProto;
 import org.apache.hadoop.yarn.federation.proto.YarnServerFederationProtos.RouterMasterKeyProto;
 import org.apache.hadoop.yarn.federation.proto.YarnServerFederationProtos.RouterMasterKeyRequestProto;
 import org.apache.hadoop.yarn.federation.proto.YarnServerFederationProtos.RouterMasterKeyResponseProto;
+import org.apache.hadoop.yarn.federation.proto.YarnServerFederationProtos.ApplicationHomeSubClusterProto;
 import org.apache.hadoop.yarn.server.federation.store.records.impl.pb.AddApplicationHomeSubClusterRequestPBImpl;
 import org.apache.hadoop.yarn.server.federation.store.records.impl.pb.AddApplicationHomeSubClusterResponsePBImpl;
 import org.apache.hadoop.yarn.server.federation.store.records.impl.pb.DeleteApplicationHomeSubClusterRequestPBImpl;
@@ -81,6 +84,8 @@
 import org.apache.hadoop.yarn.server.federation.store.records.impl.pb.RouterMasterKeyPBImpl;
 import org.apache.hadoop.yarn.server.federation.store.records.impl.pb.RouterMasterKeyRequestPBImpl;
 import org.apache.hadoop.yarn.server.federation.store.records.impl.pb.RouterMasterKeyResponsePBImpl;
+import org.apache.hadoop.yarn.server.federation.store.records.impl.pb.ApplicationHomeSubClusterPBImpl;
+import org.apache.hadoop.yarn.server.federation.store.records.impl.pb.GetReservationHomeSubClusterRequestPBImpl;
 import org.apache.hadoop.yarn.server.records.Version;
 import org.junit.BeforeClass;
 import org.junit.Test;
@@ -99,6 +104,7 @@ public static void setup() throws Exception {
     generateByNewInstance(ApplicationHomeSubCluster.class);
     generateByNewInstance(SubClusterPolicyConfiguration.class);
     generateByNewInstance(RouterMasterKey.class);
+    generateByNewInstance(ReservationId.class);
   }
 
   @Test
@@ -284,4 +290,16 @@ public void testRouterMasterKeyRequest() throws Exception {
   public void testRouterMasterKeyResponse() throws Exception {
     validatePBImplRecord(RouterMasterKeyResponsePBImpl.class, RouterMasterKeyResponseProto.class);
   }
+
+  @Test
+  public void testApplicationHomeSubCluster() throws Exception {
+    validatePBImplRecord(ApplicationHomeSubClusterPBImpl.class,
+        ApplicationHomeSubClusterProto.class);
+  }
+
+  @Test
+  public void testGetReservationHomeSubClusterRequest() throws Exception {
+    validatePBImplRecord(GetReservationHomeSubClusterRequestPBImpl.class,
+        GetReservationHomeSubClusterRequestProto.class);
+  }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/test/java/org/apache/hadoop/yarn/server/federation/utils/FederationPoliciesTestUtil.java

@@ -276,8 +276,8 @@ public static FederationStateStoreFacade initFacade(
    * @throws YarnException in case the initialization is not successful.
    */
   public static FederationStateStoreFacade initFacade() throws YarnException {
-    return initFacade(new ArrayList<>(), mock(SubClusterPolicyConfiguration
+    SubClusterPolicyConfiguration policyConfiguration =
-        .class));
+        SubClusterPolicyConfiguration.newInstance(null, null, null);
+    return initFacade(new ArrayList<>(), policyConfiguration);
   }
-
 }