YARN-5092. TestRMDelegationTokens fails intermittently. Contributed by Jason Lowe.

(cherry picked from commit 557a245d83c358211f51b3ab1950d707a70679d9)

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/security/TestRMDelegationTokens.java

@@ -43,7 +43,6 @@ import org.apache.hadoop.yarn.server.resourcemanager.TestRMRestart.TestSecurityM
 import org.apache.hadoop.yarn.server.resourcemanager.recovery.MemoryRMStateStore;
 import org.apache.hadoop.yarn.server.resourcemanager.recovery.RMStateStore;
 import org.apache.hadoop.yarn.server.resourcemanager.recovery.RMStateStore.RMState;
-import org.apache.hadoop.yarn.server.resourcemanager.scheduler.fair.FairScheduler;
 import org.apache.hadoop.yarn.util.ConverterUtils;
 import org.apache.log4j.Level;
 import org.apache.log4j.LogManager;
@@ -54,23 +53,24 @@ import org.junit.Test;
 
 public class TestRMDelegationTokens {
 
-  private YarnConfiguration conf;
+  private YarnConfiguration testConf;
 
   @Before
   public void setup() {
     Logger rootLogger = LogManager.getRootLogger();
     rootLogger.setLevel(Level.DEBUG);
     ExitUtil.disableSystemExit();
-    conf = new YarnConfiguration();
-    UserGroupInformation.setConfiguration(conf);
-    conf.set(YarnConfiguration.RM_STORE, MemoryRMStateStore.class.getName());
-    conf.set(YarnConfiguration.RM_SCHEDULER, FairScheduler.class.getName());
+    testConf = new YarnConfiguration();
+    UserGroupInformation.setLoginUser(null);
+    UserGroupInformation.setConfiguration(testConf);
+    testConf.set(YarnConfiguration.RM_STORE,
+        MemoryRMStateStore.class.getName());
   }
 
   // Test the DT mast key in the state-store when the mast key is being rolled.
   @Test(timeout = 15000)
   public void testRMDTMasterKeyStateOnRollingMasterKey() throws Exception {
-    Configuration conf = new Configuration();
+    Configuration conf = new Configuration(testConf);
     conf.set("hadoop.security.authentication", "kerberos");
     UserGroupInformation.setConfiguration(conf);
     MemoryRMStateStore memStore = new MemoryRMStateStore();
@@ -93,9 +93,6 @@ public class TestRMDelegationTokens {
         rm1.getRMContext().getRMDelegationTokenSecretManager();
     // assert all master keys are saved
     Assert.assertEquals(dtSecretManager.getAllMasterKeys(), rmDTMasterKeyState);
-    Set<DelegationKey> expiringKeys = new HashSet<DelegationKey>();
-    expiringKeys.addAll(dtSecretManager.getAllMasterKeys());
-
 
     // request to generate a RMDelegationToken
     GetDelegationTokenRequest request = mock(GetDelegationTokenRequest.class);
@@ -131,13 +128,13 @@ public class TestRMDelegationTokens {
   @Test(timeout = 15000)
   public void testRemoveExpiredMasterKeyInRMStateStore() throws Exception {
     MemoryRMStateStore memStore = new MemoryRMStateStore();
-    memStore.init(conf);
+    memStore.init(testConf);
     RMState rmState = memStore.getState();
 
     Set<DelegationKey> rmDTMasterKeyState =
         rmState.getRMDTSecretManagerState().getMasterKeyState();
 
-    MockRM rm1 = new MyMockRM(conf, memStore);
+    MockRM rm1 = new MyMockRM(testConf, memStore);
     rm1.start();
     RMDelegationTokenSecretManager dtSecretManager =
         rm1.getRMContext().getRMDelegationTokenSecretManager();
@@ -159,6 +156,7 @@ public class TestRMDelegationTokens {
         break;
       Thread.sleep(500);
     }
+    rm1.stop();
   }
 
   class MyMockRM extends TestSecurityMockRM {
@@ -169,7 +167,7 @@ public class TestRMDelegationTokens {
 
     @Override
     protected RMSecretManagerService createRMSecretManagerService() {
-      return new RMSecretManagerService(conf, rmContext) {
+      return new RMSecretManagerService(testConf, rmContext) {
 
         @Override
         protected RMDelegationTokenSecretManager
@@ -208,7 +206,17 @@ public class TestRMDelegationTokens {
       for (int keyId : allKeys.keySet()) {
         if (keyId == currentId) {
           DelegationKey currentKey = allKeys.get(keyId);
-          Assert.assertTrue(rmDTMasterKeyState.contains(currentKey));
+          // There's a small window where the key expiry has changed in memory
+          // but not the state store yet, and DelegationKey hashcode/equals
+          // uses the expiry so the contains method will fail to find it.
+          boolean found = false;
+          for (DelegationKey k : rmDTMasterKeyState) {
+            if (k.getKeyId() == keyId) {
+              found = true;
+              break;
+            }
+          }
+          Assert.assertTrue(found);
           return currentKey;
         }
       }