From 0066526dcefca4043aff1ebf73d03dd554ec4e62 Mon Sep 17 00:00:00 2001
From: Chung En Lee <wfps1210@gmail.com>
Date: Thu, 19 Dec 2024 22:02:43 +0800
Subject: [PATCH] HDDS-11869. Enable OM Ratis in
 TestOzoneDelegationTokenSecretManager (#7594)

---
 ...TestOzoneDelegationTokenSecretManager.java | 20 ++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/security/TestOzoneDelegationTokenSecretManager.java b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/security/TestOzoneDelegationTokenSecretManager.java
index c0fdb7a8c21..a9239b5639a 100644
--- a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/security/TestOzoneDelegationTokenSecretManager.java
+++ b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/security/TestOzoneDelegationTokenSecretManager.java
@@ -59,7 +59,6 @@
 import org.apache.hadoop.security.token.Token;
 import org.apache.hadoop.util.Time;
 
-import static org.apache.hadoop.ozone.om.OMConfigKeys.OZONE_OM_RATIS_ENABLE_KEY;
 import static org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.OMTokenProto.Type.S3AUTHINFO;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -132,14 +131,6 @@ public void setUp() throws Exception {
 
   private OzoneConfiguration createNewTestPath() throws IOException {
     OzoneConfiguration config = new OzoneConfiguration();
-    // When ratis is enabled, tokens are not updated to the store directly by
-    // OzoneDelegationTokenSecretManager. Tokens are updated via Ratis
-    // through the DoubleBuffer. Hence, to test
-    // OzoneDelegationTokenSecretManager, we should disable OM Ratis.
-    // TODO: Once HA and non-HA code paths are merged in
-    //  OzoneDelegationTokenSecretManager, this test should be updated to
-    //  test both ratis enabled and disabled case.
-    config.setBoolean(OZONE_OM_RATIS_ENABLE_KEY, false);
     File newFolder = folder.toFile();
     if (!newFolder.exists()) {
       assertTrue(newFolder.mkdirs());
@@ -257,6 +248,7 @@ private void testRenewTokenSuccessHelper(boolean restartSecretManager)
     Token<OzoneTokenIdentifier> token = secretManager.createToken(TEST_USER,
         TEST_USER,
         TEST_USER);
+    addToTokenStore(token);
     Thread.sleep(10 * 5);
 
     if (restartSecretManager) {
@@ -264,6 +256,7 @@ private void testRenewTokenSuccessHelper(boolean restartSecretManager)
     }
 
     long renewalTime = secretManager.renewToken(token, TEST_USER.toString());
+    addToTokenStore(token);
     assertThat(renewalTime).isGreaterThan(0);
   }
 
@@ -287,6 +280,7 @@ public void testRenewTokenFailure() throws Exception {
     secretManager.start(certificateClient);
     Token<OzoneTokenIdentifier> token = secretManager.createToken(TEST_USER,
         TEST_USER, TEST_USER);
+    addToTokenStore(token);
     AccessControlException exception =
         assertThrows(AccessControlException.class,
             () -> secretManager.renewToken(token, "rougeUser"));
@@ -354,6 +348,7 @@ public void testCancelTokenSuccess() throws Exception {
     secretManager.start(certificateClient);
     Token<OzoneTokenIdentifier> token = secretManager.createToken(TEST_USER,
         TEST_USER, TEST_USER);
+    addToTokenStore(token);
     secretManager.cancelToken(token, TEST_USER.toString());
   }
 
@@ -512,4 +507,11 @@ private void validateHash(byte[] hash, byte[] identifier) throws Exception {
         .setSecretKeyClient(secretKeyClient)
         .build();
   }
+
+  private void addToTokenStore(Token<OzoneTokenIdentifier> token) throws IOException {
+    OzoneTokenIdentifier ozoneTokenIdentifier = OzoneTokenIdentifier.
+        readProtoBuf(token.getIdentifier());
+    long renewDate = secretManager.updateToken(token, ozoneTokenIdentifier, expiryTime);
+    om.getMetadataManager().getDelegationTokenTable().put(ozoneTokenIdentifier, renewDate);
+  }
 }