From 7d998cb6278ad5963ef4e2a384b81ee19ff0a489 Mon Sep 17 00:00:00 2001
From: Abhishek Modi <abmodi@apache.org>
Date: Sat, 31 Aug 2019 10:03:09 +0530
Subject: [PATCH] YARN-9800. TestRMDelegationTokens can fail in
 testRemoveExpiredMasterKeyInRMStateStore. Contributed by Adam Antal.

---
 .../security/TestRMDelegationTokens.java      | 33 +++++++++++--------
 1 file changed, 19 insertions(+), 14 deletions(-)

diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/security/TestRMDelegationTokens.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/security/TestRMDelegationTokens.java
index c660f9b71b3..94bf3edfdd1 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/security/TestRMDelegationTokens.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/security/TestRMDelegationTokens.java
@@ -69,6 +69,22 @@ public void setup() {
     UserGroupInformation.setConfiguration(testConf);
   }
 
+  private static void assertMasterKeysAreSaved(
+      Set<DelegationKey> rmDTMasterKeyState,
+      RMDelegationTokenSecretManager dtSecretManager) {
+    dtSecretManager.getAllMasterKeys().forEach(managerKey -> {
+      int keyId = managerKey.getKeyId();
+      boolean found = false;
+      for (DelegationKey stateKey: rmDTMasterKeyState) {
+        if (stateKey.getKeyId() == keyId) {
+          found = true;
+          break;
+        }
+      }
+      Assert.assertTrue("Master key not found: " + keyId, found);
+    });
+  }
+
   // Test the DT mast key in the state-store when the mast key is being rolled.
   @Test(timeout = 15000)
   public void testRMDTMasterKeyStateOnRollingMasterKey() throws Exception {
@@ -96,18 +112,8 @@ public void testRMDTMasterKeyStateOnRollingMasterKey() throws Exception {
 
     RMDelegationTokenSecretManager dtSecretManager =
         rm1.getRMContext().getRMDelegationTokenSecretManager();
-    // assert all master keys are saved
-    dtSecretManager.getAllMasterKeys().forEach(managerKey -> {
-      int keyId = managerKey.getKeyId();
-      boolean found = false;
-      for (DelegationKey stateKey: rmDTMasterKeyState) {
-        if (stateKey.getKeyId() == keyId) {
-          found = true;
-          break;
-        }
-      }
-      Assert.assertTrue("Master key not found: " + keyId, found);
-    });
+
+    assertMasterKeysAreSaved(rmDTMasterKeyState, dtSecretManager);
 
     // request to generate a RMDelegationToken
     GetDelegationTokenRequest request = mock(GetDelegationTokenRequest.class);
@@ -154,8 +160,7 @@ public void testRemoveExpiredMasterKeyInRMStateStore() throws Exception {
     RMDelegationTokenSecretManager dtSecretManager =
         rm1.getRMContext().getRMDelegationTokenSecretManager();
 
-    // assert all master keys are saved
-    Assert.assertEquals(dtSecretManager.getAllMasterKeys(), rmDTMasterKeyState);
+    assertMasterKeysAreSaved(rmDTMasterKeyState, dtSecretManager);
     Set<DelegationKey> expiringKeys = new HashSet<DelegationKey>();
     expiringKeys.addAll(dtSecretManager.getAllMasterKeys());