fix UTs

Signed-off-by: Jing Zhang <[email protected]>
opensearch-project · Oct 28, 2024 · d424ff0 · d424ff0
1 parent e51d00e
commit d424ff0
Showing 1 changed file with 4 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
@@ -462,8 +462,8 @@ public void decrypt() {
 
     @Test
     public void encrypt_NullMasterKey_NullMasterKey_MasterKeyNotExistInIndex() {
-        exceptionRule.expect(ResourceNotFoundException.class);
-        exceptionRule.expectMessage(MASTER_KEY_NOT_READY_ERROR);
+        exceptionRule.expect(MLException.class);
+        exceptionRule.expectMessage("Fetching master key timed out.");
 
         doAnswer(invocation -> {
             ActionListener<GetResponse> listener = invocation.getArgument(1);
@@ -517,8 +517,8 @@ public void decrypt_NoResponseToInitConfigIndex() {
 
     @Test
     public void decrypt_MLConfigIndexNotFound() {
-        exceptionRule.expect(ResourceNotFoundException.class);
-        exceptionRule.expectMessage(MASTER_KEY_NOT_READY_ERROR);
+        exceptionRule.expect(MLException.class);
+        exceptionRule.expectMessage("Fetching master key timed out.");
 
         Metadata metadata = new Metadata.Builder().indices(ImmutableMap.of()).build();
         when(clusterState.metadata()).thenReturn(metadata);