From 4b5aa9aac272d400952dfdd38d158912bd9ce269 Mon Sep 17 00:00:00 2001 From: Jing Zhang Date: Wed, 13 Nov 2024 13:13:43 -0800 Subject: [PATCH] fix UT Signed-off-by: Jing Zhang --- .../ml/engine/encryptor/EncryptorImplTest.java | 10 ++++++---- 1 file changed, 6 insertions(+), 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 index 25a3e2d49d..211ea017c3 100644 --- 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 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);