fix UT

Signed-off-by: Jing Zhang <[email protected]>
opensearch-project · Nov 13, 2024 · 5f29b5b · 5f29b5b
1 parent c1d09da
commit 5f29b5b
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/ml-algorithms/src/test/java/org/opensearch/ml/engine/encryptor/EncryptorImplTest.java b/ml-algorithms/src/test/java/org/opensearch/ml/engine/encryptor/EncryptorImplTest.java
@@ -7,6 +7,7 @@
 import static org.opensearch.ml.common.CommonValue.CREATE_TIME_FIELD;
 import static org.opensearch.ml.common.CommonValue.MASTER_KEY;
 import static org.opensearch.ml.common.CommonValue.ML_CONFIG_INDEX;
+import static org.opensearch.ml.engine.encryptor.EncryptorImpl.MASTER_KEY_NOT_READY_ERROR;
 
 import java.time.Instant;
 
@@ -17,6 +18,7 @@
 import org.junit.rules.ExpectedException;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.opensearch.ResourceNotFoundException;
 import org.opensearch.Version;
 import org.opensearch.action.get.GetResponse;
 import org.opensearch.client.Client;
@@ -129,8 +131,8 @@ public void decrypt() {
 
     @Test
     public void encrypt_NullMasterKey_NullMasterKey_MasterKeyNotExistInIndex() {
-        exceptionRule.expect(MLException.class);
-        exceptionRule.expectMessage("Fetching master key timed out.");
+        exceptionRule.expect(ResourceNotFoundException.class);
+        exceptionRule.expectMessage(MASTER_KEY_NOT_READY_ERROR);
 
         doAnswer(invocation -> {
             ActionListener<GetResponse> listener = invocation.getArgument(1);
@@ -163,8 +165,8 @@ public void decrypt_NullMasterKey_GetMasterKey_Exception() {
 
     @Test
     public void decrypt_MLConfigIndexNotFound() {
-        exceptionRule.expect(MLException.class);
-        exceptionRule.expectMessage("Fetching master key timed out.");
+        exceptionRule.expect(ResourceNotFoundException.class);
+        exceptionRule.expectMessage(MASTER_KEY_NOT_READY_ERROR);
 
         Metadata metadata = new Metadata.Builder().indices(ImmutableMap.of()).build();
         when(clusterState.metadata()).thenReturn(metadata);