Add test cases specific to handling secret manager in a distributed setup

hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/security/token/delegation/AbstractDelegationTokenSecretManager.java

@@ -237,7 +237,7 @@ extends AbstractDelegationTokenIdentifier>
 
   // RM
   protected void removeStoredToken(TokenIdent ident) throws IOException {
-
+    return;
   }
   // RM
   protected void updateStoredToken(TokenIdent ident, long renewDate) throws IOException {
@@ -386,6 +386,15 @@ extends AbstractDelegationTokenIdentifier>
     updateStoredToken(ident, tokenInfo.getRenewDate());
   }
 
+  protected void removeToken(TokenIdent ident) throws IOException {
+    DelegationTokenInformation info = currentTokens.remove(ident);
+    if (info == null) {
+      throw new InvalidToken("Token not found " + formatTokenId(ident));
+    }
+    removeTokenForOwnerStats(ident);
+    removeStoredToken(ident);
+  }
+
   /**
    * This method is intended to be used for recovering persisted delegation
    * tokens. Tokens that have an unknown <code>DelegationKey</code> are
@@ -691,13 +700,8 @@ extends AbstractDelegationTokenIdentifier>
       throw new AccessControlException(canceller
           + " is not authorized to cancel the token " + formatTokenId(id));
     }
-    DelegationTokenInformation info = currentTokens.remove(id);
-    if (info == null) {
-      throw new InvalidToken("Token not found " + formatTokenId(id));
-    }
     METRICS.trackRemoveToken(() -> {
-      removeTokenForOwnerStats(id);
+      removeToken(id);
-      removeStoredToken(id);
     });
     return id;
   }

hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/security/token/delegation/SQLDelegationTokenSecretManager.java

@@ -161,7 +161,7 @@ public abstract class SQLDelegationTokenSecretManager<TokenIdent
    *         null if it doesn't exist in the database.
    */
   @Override
-  protected DelegationTokenInformation getTokenInfo(TokenIdent ident) {
+  protected DelegationTokenInformation getTokenInfo(TokenIdent ident) throws IOException {
     // Look for token in local cache
     DelegationTokenInformation tokenInfo = super.getTokenInfo(ident);
 
@@ -302,7 +302,7 @@ public abstract class SQLDelegationTokenSecretManager<TokenIdent
    *         if it doesn't exist in the database.
    */
   @Override
-  protected DelegationKey getDelegationKey(int keyId) {
+  protected DelegationKey getDelegationKey(int keyId) throws IOException {
     // Look for delegation key in local cache
     DelegationKey delegationKey = super.getDelegationKey(keyId);
 

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

@@ -138,10 +138,7 @@ public class MemoryFederationStateStore implements FederationStateStore {
 
   @Override
   public void close() {
-    membership = null;
+    init(new Configuration());
-    applications = null;
-    reservations = null;
-    policies = null;
   }
 
   @Override
@@ -464,13 +461,13 @@ public class MemoryFederationStateStore implements FederationStateStore {
     RouterMasterKey masterKey = request.getRouterMasterKey();
     DelegationKey delegationKey = getDelegationKeyByMasterKey(masterKey);
 
-    Set<DelegationKey> rmDTMasterKeyState = routerRMSecretManagerState.getMasterKeyState();
+    Map<Integer, DelegationKey> rmDTMasterKeyState = routerRMSecretManagerState.getMasterKeyState();
-    if (rmDTMasterKeyState.contains(delegationKey)) {
+    if (rmDTMasterKeyState.containsKey(delegationKey.getKeyId())) {
       FederationStateStoreUtils.logAndThrowStoreException(LOG,
           "Error storing info for RMDTMasterKey with keyID: %s.", delegationKey.getKeyId());
     }
 
-    routerRMSecretManagerState.getMasterKeyState().add(delegationKey);
+    routerRMSecretManagerState.getMasterKeyState().put(delegationKey.getKeyId(), delegationKey);
     LOG.info("Store Router-RMDT master key with key id: {}. Currently rmDTMasterKeyState size: {}",
         delegationKey.getKeyId(), rmDTMasterKeyState.size());
 
@@ -485,8 +482,8 @@ public class MemoryFederationStateStore implements FederationStateStore {
     DelegationKey delegationKey = getDelegationKeyByMasterKey(masterKey);
 
     LOG.info("Remove Router-RMDT master key with key id: {}.", delegationKey.getKeyId());
-    Set<DelegationKey> rmDTMasterKeyState = routerRMSecretManagerState.getMasterKeyState();
+    Map<Integer, DelegationKey> rmDTMasterKeyState = routerRMSecretManagerState.getMasterKeyState();
-    rmDTMasterKeyState.remove(delegationKey);
+    rmDTMasterKeyState.remove(delegationKey.getKeyId());
 
     return RouterMasterKeyResponse.newInstance(masterKey);
   }
@@ -494,12 +491,10 @@ public class MemoryFederationStateStore implements FederationStateStore {
   @Override
   public RouterMasterKeyResponse getMasterKeyByDelegationKey(RouterMasterKeyRequest request)
       throws YarnException, IOException {
-    // Restore the DelegationKey from the request
     RouterMasterKey masterKey = request.getRouterMasterKey();
-    DelegationKey delegationKey = getDelegationKeyByMasterKey(masterKey);
+    Map<Integer, DelegationKey> rmDTMasterKeyState = routerRMSecretManagerState.getMasterKeyState();
-
+    DelegationKey delegationKey = rmDTMasterKeyState.get(masterKey.getKeyId());
-    Set<DelegationKey> rmDTMasterKeyState = routerRMSecretManagerState.getMasterKeyState();
+    if (delegationKey == null) {
-    if (!rmDTMasterKeyState.contains(delegationKey)) {
       throw new IOException("GetMasterKey with keyID: " + masterKey.getKeyId() +
           " does not exist.");
     }

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

@@ -30,7 +30,7 @@ public class RouterRMDTSecretManagerState {
   // DTIdentifier -> renewDate
   private Map<Integer, RouterStoreToken> delegationTokenState = new HashMap<>();
 
-  private Set<DelegationKey> masterKeyState = new HashSet<>();
+  private Map<Integer, DelegationKey> masterKeyState = new HashMap<>();
 
   private int dtSequenceNumber = 0;
 
@@ -38,7 +38,7 @@ public class RouterRMDTSecretManagerState {
     return delegationTokenState;
   }
 
-  public Set<DelegationKey> getMasterKeyState() {
+  public Map<Integer, DelegationKey> getMasterKeyState() {
     return masterKeyState;
   }
 

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/test/java/org/apache/hadoop/yarn/server/federation/store/impl/TestMemoryFederationStateStore.java

@@ -59,10 +59,10 @@ public class TestMemoryFederationStateStore extends FederationStateStoreBaseTest
         memoryStateStore.getRouterRMSecretManagerState();
     assertNotNull(secretManagerState);
 
-    Set<DelegationKey> delegationKeys = secretManagerState.getMasterKeyState();
+    Map<Integer, DelegationKey> delegationKeys = secretManagerState.getMasterKeyState();
     assertNotNull(delegationKeys);
 
-    assertTrue(delegationKeys.contains(delegationKey));
+    assertTrue(delegationKeys.containsKey(delegationKey.getKeyId()));
 
     RouterMasterKey resultRouterMasterKey = RouterMasterKey.newInstance(delegationKey.getKeyId(),
         ByteBuffer.wrap(delegationKey.getEncodedKey()), delegationKey.getExpiryDate());

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-router/src/main/java/org/apache/hadoop/yarn/server/router/security/RouterDelegationTokenSecretManager.java

@@ -145,7 +145,7 @@ public class RouterDelegationTokenSecretManager
   }
 
   /**
-   * no-op as expiry of stored tokens is upto the state store in a stateless secret manager
+   * no-op as removal of tokens is handled in removeToken()
    */
   @Override
   public void removeStoredToken(RMDelegationTokenIdentifier identifier) {
@@ -188,7 +188,6 @@ public class RouterDelegationTokenSecretManager
   @Override
   protected void storeToken(RMDelegationTokenIdentifier rmDelegationTokenIdentifier,
       DelegationTokenInformation tokenInfo) throws IOException {
-    this.addTokenForOwnerStats(rmDelegationTokenIdentifier);
     try {
       storeNewToken(rmDelegationTokenIdentifier, tokenInfo);
     } catch (YarnException e) {
@@ -208,6 +207,16 @@ public class RouterDelegationTokenSecretManager
     }
   }
 
+  @Override
+  protected void removeToken(RMDelegationTokenIdentifier identifier) throws IOException {
+    try {
+      federationFacade.removeStoredToken(identifier);
+    } catch (YarnException e) {
+      e.printStackTrace();
+      throw new IOException(e); // Wrap YarnException as an IOException to adhere to removeToken contract
+    }
+  }
+
   @Override
   protected DelegationTokenInformation getTokenInfo(RMDelegationTokenIdentifier ident) throws IOException {
       try {

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-router/src/test/java/org/apache/hadoop/yarn/server/router/security/TestRouterDelegationTokenSecretManager.java

@@ -18,52 +18,55 @@
 package org.apache.hadoop.yarn.server.router.security;
 
 import org.apache.hadoop.io.Text;
+import org.apache.hadoop.security.token.SecretManager;
 import org.apache.hadoop.security.token.Token;
-import org.apache.hadoop.security.token.delegation.AbstractDelegationTokenSecretManager;
-import org.apache.hadoop.security.token.delegation.DelegationKey;
-import org.apache.hadoop.test.LambdaTestUtils;
-import org.apache.hadoop.util.Time;
 import org.apache.hadoop.yarn.security.client.RMDelegationTokenIdentifier;
-import org.apache.hadoop.yarn.server.router.clientrm.RouterClientRMService;
+import org.apache.hadoop.yarn.server.federation.utils.FederationStateStoreFacade;
+import org.junit.After;
 import org.junit.Assert;
 import org.junit.Before;
+import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.jupiter.api.Assertions;
-import org.slf4j.Logger;
+import org.mockito.Mockito;
-import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
 
-import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertArrayEquals;
-
 public class TestRouterDelegationTokenSecretManager extends AbstractSecureRouterTest {
 
-  private static final Logger LOG =
+  private volatile RouterDelegationTokenSecretManager secretManager_1;
-      LoggerFactory.getLogger(TestRouterDelegationTokenSecretManager.class);
+  private volatile RouterDelegationTokenSecretManager secretManager_2;
-
+  private final Text owner = new Text("owner");
-  private RouterDelegationTokenSecretManager secretManager_1;
+  private final Text renewer = new Text("renewer");
-  private RouterDelegationTokenSecretManager secretManager_2;
+  private final Text realUser = new Text("realUser");
-  private final Text owner = new Text("hadoop");
+  private final int keyUpdateInterval = 1000;
-  private final Text renewer = new Text("yarn");
+  private final int tokenRenewInterval = 2000;
-  private final Text realUser = new Text("router");
+  private final int tokenMaxLifeTime = 10000;
 
   @Before
   public void setup() {
 
     // Setup multiple secret managers to validate stateless secret managers.
     // They are using same instance of FederationStateStoreFacade thus the in memory state store is shared
-    secretManager_1 = new RouterDelegationTokenSecretManager(
+    secretManager_1 = Mockito.spy(new RouterDelegationTokenSecretManager(
-        1000, 10000, 1000, 100
+        keyUpdateInterval, tokenMaxLifeTime, tokenRenewInterval, 100)
     );
-    secretManager_2 = new RouterDelegationTokenSecretManager(
+    secretManager_2 = Mockito.spy(new RouterDelegationTokenSecretManager(
-        1000, 10000, 1000, 100
+        keyUpdateInterval, tokenMaxLifeTime, tokenRenewInterval, 100)
     );
   }
 
+  @After
+  public void cleanup() throws Exception {
+    secretManager_1.stopThreads();
+    secretManager_2.stopThreads();
+    secretManager_1 = null;
+    secretManager_2 = null;
+    FederationStateStoreFacade.getInstance().getStateStore().close();
+  }
+
   @Test
-  public void testNewTokenVerification() throws IOException {
+  public void testNewTokenIsVerifiedAcrossManagers() throws IOException {
 
     secretManager_1.startThreads();
     RMDelegationTokenIdentifier tokenIdentifier = new RMDelegationTokenIdentifier(owner, renewer, realUser);
@@ -77,162 +80,172 @@ public class TestRouterDelegationTokenSecretManager extends AbstractSecureRouter
 
     secretManager_2.startThreads();
     RMDelegationTokenIdentifier tokenIdentifier_3 = secretManager_2.decodeTokenIdentifier(token2);
-    Assertions.assertDoesNotThrow(() -> secretManager_2.verifyToken(tokenIdentifier_3, token.getPassword()));
+    Assertions.assertDoesNotThrow(() -> secretManager_2.verifyToken(tokenIdentifier_3, token2.getPassword()));
   }
 
   @Test
-  public void testRouterStoreNewMasterKey() throws Exception {
+  public void testMasterKeyIsRolled() throws IOException, InterruptedException {
-    LOG.info("Test RouterDelegationTokenSecretManager: StoreNewMasterKey.");
 
-    // Start the Router in Secure Mode
+    secretManager_1.startThreads();
-    startSecureRouter();
+    RMDelegationTokenIdentifier tokenIdentifier1 = new RMDelegationTokenIdentifier(owner, renewer, realUser);
+    new Token<>(tokenIdentifier1, secretManager_1);
 
-    // Store NewMasterKey
+    RMDelegationTokenIdentifier tokenIdentifier2 = new RMDelegationTokenIdentifier(owner, renewer, realUser);
-    RouterClientRMService routerClientRMService = this.getRouter().getClientRMProxyService();
+    new Token<>(tokenIdentifier2, secretManager_1);
-    RouterDelegationTokenSecretManager secretManager =
-        routerClientRMService.getRouterDTSecretManager();
-    DelegationKey storeKey = new DelegationKey(1234, 4321, "keyBytes".getBytes());
-    secretManager.storeNewMasterKey(storeKey);
 
-    // Get DelegationKey
+    // Check multiple tokens have same master key
-    DelegationKey responseKey = secretManager.getDelegationKey(1234);
+    Assert.assertEquals(tokenIdentifier1.getMasterKeyId(), tokenIdentifier2.getMasterKeyId());
+    // Sleep until master key is updated
+    Thread.sleep(keyUpdateInterval + 100);
 
-    assertNotNull(responseKey);
+    RMDelegationTokenIdentifier tokenIdentifier3 = new RMDelegationTokenIdentifier(owner, renewer, realUser);
-    assertEquals(storeKey.getExpiryDate(), responseKey.getExpiryDate());
+    new Token<>(tokenIdentifier3, secretManager_1);
-    assertEquals(storeKey.getKeyId(), responseKey.getKeyId());
+    // Check master key is updated
-    assertArrayEquals(storeKey.getEncodedKey(), responseKey.getEncodedKey());
+    Assert.assertNotEquals(tokenIdentifier1.getMasterKeyId(), tokenIdentifier3.getMasterKeyId());
-    assertEquals(storeKey, responseKey);
 
-    stopSecureRouter();
   }
 
   @Test
-  public void testRouterRemoveStoredMasterKey() throws Exception {
+  public void testNewTokenIsCancelledAcrossManagers() throws IOException {
-    LOG.info("Test RouterDelegationTokenSecretManager: RemoveStoredMasterKey.");
 
-    // Start the Router in Secure Mode
+    secretManager_1.startThreads();
-    startSecureRouter();
+    RMDelegationTokenIdentifier tokenIdentifier = new RMDelegationTokenIdentifier(owner, renewer, realUser);
+    Token<RMDelegationTokenIdentifier> token = new Token<>(tokenIdentifier, secretManager_1);
 
-    // Store NewMasterKey
+    Token<RMDelegationTokenIdentifier> token2 = new Token<>();
-    RouterClientRMService routerClientRMService = this.getRouter().getClientRMProxyService();
+    token2.decodeFromUrlString(token.encodeToUrlString());
-    RouterDelegationTokenSecretManager secretManager =
-        routerClientRMService.getRouterDTSecretManager();
-    DelegationKey storeKey = new DelegationKey(1234, 4321, "keyBytes".getBytes());
-    secretManager.storeNewMasterKey(storeKey);
 
-    // Remove DelegationKey
+    secretManager_2.startThreads();
-    secretManager.removeStoredMasterKey(storeKey);
+    secretManager_2.cancelToken(token2, owner.toString());
 
-    // Get DelegationKey
+    RMDelegationTokenIdentifier tokenIdentifier_2 = secretManager_1.decodeTokenIdentifier(token2);
-    LambdaTestUtils.intercept(IOException.class,
+    Assertions.assertThrows(SecretManager.InvalidToken.class,
-        "GetMasterKey with keyID: " + storeKey.getKeyId() + " does not exist.",
+        () -> secretManager_1.verifyToken(tokenIdentifier_2, token2.getPassword())
-        () -> secretManager.getDelegationKey(1234));
+    );
 
-    stopSecureRouter();
   }
 
   @Test
-  public void testRouterStoreNewToken() throws Exception {
+  public void testNewTokenIsRenewedAcrossManagers() throws IOException, InterruptedException {
-    LOG.info("Test RouterDelegationTokenSecretManager: StoreNewToken.");
 
-    // Start the Router in Secure Mode
+    secretManager_1.startThreads();
-    startSecureRouter();
+    RMDelegationTokenIdentifier tokenIdentifier = new RMDelegationTokenIdentifier(owner, renewer, realUser);
+    Token<RMDelegationTokenIdentifier> token = new Token<>(tokenIdentifier, secretManager_1);
 
-    // Store new rm-token
+    Token<RMDelegationTokenIdentifier> token2 = new Token<>();
-    RouterClientRMService routerClientRMService = this.getRouter().getClientRMProxyService();
+    token2.decodeFromUrlString(token.encodeToUrlString());
-    RouterDelegationTokenSecretManager secretManager =
-        routerClientRMService.getRouterDTSecretManager();
-    RMDelegationTokenIdentifier dtId1 = new RMDelegationTokenIdentifier(
-        new Text("owner1"), new Text("renewer1"), new Text("realuser1"));
-    int sequenceNumber = 1;
-    dtId1.setSequenceNumber(sequenceNumber);
-    long renewDate1 = Time.now();
-    secretManager.storeNewToken(dtId1, renewDate1);
 
-    // query rm-token
+    Thread.sleep(tokenRenewInterval / 2 + 100);
-    RMDelegationTokenIdentifier dtId2 = new RMDelegationTokenIdentifier(
+    secretManager_2.startThreads();
-        new Text("owner1"), new Text("renewer1"), new Text("realuser1"));
+    secretManager_2.renewToken(token2, renewer.toString());
-    dtId2.setSequenceNumber(sequenceNumber);
-    AbstractDelegationTokenSecretManager.DelegationTokenInformation dtId3 = secretManager.getTokenInfo(dtId2);
-    Assert.assertEquals(renewDate1, dtId3.getRenewDate());
 
-    // query rm-token2 not exists
+    Thread.sleep(tokenRenewInterval / 2 + 100);
-    sequenceNumber++;
+    RMDelegationTokenIdentifier tokenIdentifier_2 = secretManager_1.decodeTokenIdentifier(token2);
-    dtId2.setSequenceNumber(sequenceNumber);
+    Assertions.assertDoesNotThrow(() -> secretManager_1.verifyToken(tokenIdentifier_2, token2.getPassword()));
-    LambdaTestUtils.intercept(IOException.class,
-        "RMDelegationToken: " + dtId2 + " does not exist.",
-        () -> secretManager.getTokenInfo(dtId2));
 
-    stopSecureRouter();
   }
 
   @Test
-  public void testRouterUpdateNewToken() throws Exception {
+  public void testTokenOperationsOnMasterKeyRollover() throws IOException, InterruptedException {
-    LOG.info("Test RouterDelegationTokenSecretManager: UpdateNewToken.");
 
-    // Start the Router in Secure Mode
+    secretManager_1.startThreads();
-    startSecureRouter();
+    RMDelegationTokenIdentifier tokenIdentifier1 = new RMDelegationTokenIdentifier(owner, renewer, realUser);
+    Token<RMDelegationTokenIdentifier> token1 = new Token<>(tokenIdentifier1, secretManager_1);
 
-    // Store new rm-token
+    // Sleep until master key is updated
-    RouterClientRMService routerClientRMService = this.getRouter().getClientRMProxyService();
+    Thread.sleep(keyUpdateInterval + 100);
-    RouterDelegationTokenSecretManager secretManager =
-        routerClientRMService.getRouterDTSecretManager();
-    RMDelegationTokenIdentifier dtId1 = new RMDelegationTokenIdentifier(
-        new Text("owner1"), new Text("renewer1"), new Text("realuser1"));
-    int sequenceNumber = 1;
-    dtId1.setSequenceNumber(sequenceNumber);
-    Long renewDate1 = Time.now();
-    secretManager.storeNewToken(dtId1, renewDate1);
 
-    sequenceNumber++;
+    RMDelegationTokenIdentifier tokenIdentifier2 = new RMDelegationTokenIdentifier(owner, renewer, realUser);
-    dtId1.setSequenceNumber(sequenceNumber);
+    new Token<>(tokenIdentifier2, secretManager_1);
-    secretManager.updateStoredToken(dtId1, renewDate1);
+    // Check master key is updated
+    Assert.assertNotEquals(tokenIdentifier1.getMasterKeyId(), tokenIdentifier2.getMasterKeyId());
 
-    // query rm-token
+    // Verify token with old master key is still considered valid
-    RMDelegationTokenIdentifier dtId2 = new RMDelegationTokenIdentifier(
+    Assertions.assertDoesNotThrow(() -> secretManager_1.verifyToken(tokenIdentifier1, token1.getPassword()));
-        new Text("owner1"), new Text("renewer1"), new Text("realuser1"));
+    // Verify token with old master key can be renewed
-    dtId2.setSequenceNumber(sequenceNumber);
+    Assertions.assertDoesNotThrow(() -> secretManager_1.renewToken(token1, renewer.toString()));
-//    RMDelegationTokenIdentifier dtId3 = secretManager.getTokenInfo(dtId2);
+    // Verify token with old master key can be cancelled
-//    assertNotNull(dtId3);
+    Assertions.assertDoesNotThrow(() -> secretManager_1.cancelToken(token1, owner.toString()));
-//    assertEquals(dtId1.getKind(), dtId3.getKind());
+    // Verify token with old master key is now cancelled
-//    assertEquals(dtId1.getOwner(), dtId3.getOwner());
+    Assert.assertThrows(SecretManager.InvalidToken.class,
-//    assertEquals(dtId1.getRealUser(), dtId3.getRealUser());
+        () -> secretManager_1.verifyToken(tokenIdentifier1, token1.getPassword()));
-//    assertEquals(dtId1.getRenewer(), dtId3.getRenewer());
-//    assertEquals(dtId1.getIssueDate(), dtId3.getIssueDate());
-//    assertEquals(dtId1.getMasterKeyId(), dtId3.getMasterKeyId());
-//    assertEquals(dtId1.getSequenceNumber(), dtId3.getSequenceNumber());
-//    assertEquals(sequenceNumber, dtId3.getSequenceNumber());
-//    assertEquals(dtId1, dtId3);
 
-    stopSecureRouter();
   }
 
   @Test
-  public void testRouterRemoveToken() throws Exception {
+  public void testMasterKeyIsNotRolledOver() throws IOException, InterruptedException {
-    LOG.info("Test RouterDelegationTokenSecretManager: RouterRemoveToken.");
 
-    // Start the Router in Secure Mode
+    secretManager_1.startThreads();
-    startSecureRouter();
+    RMDelegationTokenIdentifier tokenIdentifier1 = new RMDelegationTokenIdentifier(owner, renewer, realUser);
+    new Token<>(tokenIdentifier1, secretManager_1);
 
-    // Store new rm-token
+    Mockito.doThrow(new IOException("failure")).when(secretManager_1).storeNewMasterKey(Mockito.any());
-    RouterClientRMService routerClientRMService = this.getRouter().getClientRMProxyService();
-    RouterDelegationTokenSecretManager secretManager =
-        routerClientRMService.getRouterDTSecretManager();
-    RMDelegationTokenIdentifier dtId1 = new RMDelegationTokenIdentifier(
-        new Text("owner1"), new Text("renewer1"), new Text("realuser1"));
-    int sequenceNumber = 1;
-    dtId1.setSequenceNumber(sequenceNumber);
-    Long renewDate1 = Time.now();
-    secretManager.storeNewToken(dtId1, renewDate1);
 
-    // Remove rm-token
+    // Sleep until master key is updated
-    secretManager.removeStoredToken(dtId1);
+    Thread.sleep(keyUpdateInterval + 100);
 
-    // query rm-token
+    RMDelegationTokenIdentifier tokenIdentifier2 = new RMDelegationTokenIdentifier(owner, renewer, realUser);
-    LambdaTestUtils.intercept(IOException.class,
+    new Token<>(tokenIdentifier2, secretManager_1);
-        "RMDelegationToken: " + dtId1 + " does not exist.",
+    // Verify master key is not updated
-        () -> secretManager.getTokenInfo(dtId1));
+    Assert.assertEquals(tokenIdentifier1.getMasterKeyId(), tokenIdentifier2.getMasterKeyId());
-
-    stopSecureRouter();
   }
+
+  @Test
+  public void testNewTokenFailsOnDBFailure() throws IOException {
+    secretManager_1.startThreads();
+    RMDelegationTokenIdentifier tokenIdentifier1 = new RMDelegationTokenIdentifier(owner, renewer, realUser);
+
+    Mockito.doThrow(new IOException("failure")).when(secretManager_1).storeToken(Mockito.any(), Mockito.any());
+    Assert.assertThrows(RuntimeException.class, () -> new Token<>(tokenIdentifier1, secretManager_1));
+  }
+
+  @Test
+  public void testTokenIsNotRenewedOnDBFailure() throws IOException, InterruptedException {
+    secretManager_1.startThreads();
+    RMDelegationTokenIdentifier tokenIdentifier1 = new RMDelegationTokenIdentifier(owner, renewer, realUser);
+
+    Token<RMDelegationTokenIdentifier> token = new Token<>(tokenIdentifier1, secretManager_1);
+    Mockito.doThrow(new IOException("failure")).when(secretManager_1).updateToken(Mockito.any(), Mockito.any());
+
+    Thread.sleep(tokenRenewInterval / 2 + 100);
+    Assert.assertThrows(IOException.class, () -> secretManager_1.renewToken(token, renewer.toString()));
+    // validate that token is currently valid
+    Assertions.assertDoesNotThrow(() -> secretManager_1.verifyToken(tokenIdentifier1, token.getPassword()));
+
+    Thread.sleep(tokenRenewInterval / 2 + 100);
+    // token is no longer valid because token renewal had failed
+    Assertions.assertThrows(SecretManager.InvalidToken.class,
+        () -> secretManager_1.verifyToken(tokenIdentifier1, token.getPassword())
+    );
+  }
+
+  @Ignore
+  public void testNewTokenFailureIfMasterKeyNotRolledOverAtAll() throws IOException, InterruptedException {
+    secretManager_1.startThreads();
+
+    // Token generation succeeds initially because master key generated on initialisation was saved
+    RMDelegationTokenIdentifier tokenIdentifier1 = new RMDelegationTokenIdentifier(owner, renewer, realUser);
+    new Token<>(tokenIdentifier1, secretManager_1);
+
+    Mockito.doThrow(new IOException("failure")).when(secretManager_1).storeNewMasterKey(Mockito.any());
+
+    // Sleep until current master key expires. New master key isn't generated because rollovers are failing
+    Thread.sleep(tokenMaxLifeTime + keyUpdateInterval + 100);
+
+    RMDelegationTokenIdentifier tokenIdentifier2 = new RMDelegationTokenIdentifier(owner, renewer, realUser);
+    // New token generation fails because current master key is expired
+    Assert.assertThrows(RuntimeException.class, () -> new Token<>(tokenIdentifier2, secretManager_1));
+  }
+
+  @Test
+  public void testMasterKeyCreationFailureOnStartup() throws IOException {
+    Mockito.doThrow(new IOException("failure")).when(secretManager_1).storeNewMasterKey(Mockito.any());
+
+    Assert.assertThrows(IOException.class, () -> secretManager_1.startThreads());
+
+    RMDelegationTokenIdentifier tokenIdentifier = new RMDelegationTokenIdentifier(owner, renewer, realUser);
+    // New token generation fails because master key is not yet set
+    Assert.assertThrows(NullPointerException.class, () -> new Token<>(tokenIdentifier, secretManager_1));
+  }
+
 }