Skip to content

Commit

Permalink
fix UT
Browse files Browse the repository at this point in the history 
Signed-off-by: Yaliang Wu <[email protected]>
  • Loading branch information
@ylwu-amzn
ylwu-amzn committed Jun 14, 2024
1 parent 18593a3 commit 00eb3a1
Show file tree
Hide file tree
Showing 2 changed files with 199 additions and 9 deletions.
9 changes: 6 additions & 3 deletions ml-algorithms/src/main/java/org/opensearch/ml/engine/encryptor/EncryptorImpl.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -132,12 +132,15 @@ private void initMasterKey() {
if (ExceptionUtils.getRootCause(e) instanceof VersionConflictEngineException) {
GetRequest getMasterKeyRequest = new GetRequest(ML_CONFIG_INDEX).id(MASTER_KEY);
try (ThreadContext.StoredContext threadContext = client.threadPool().getThreadContext().stashContext()) {
client.get(getMasterKeyRequest, ActionListener.wrap(getMasterKey -> {
if (getMasterKey.isExists()) {
final String masterKey = (String) getMasterKey.getSourceAsMap().get(MASTER_KEY);
client.get(getMasterKeyRequest, ActionListener.wrap(getMasterKeyResponse -> {
if (getMasterKeyResponse != null && getMasterKeyResponse.isExists()) {
final String masterKey = (String) getMasterKeyResponse.getSourceAsMap().get(MASTER_KEY);
this.masterKey = masterKey;
log.info("ML encryption master key already initialized, no action needed");
latch.countDown();
} else {
exceptionRef.set(new ResourceNotFoundException(MASTER_KEY_NOT_READY_ERROR));
latch.countDown();
}
}, error -> {
log.debug("Failed to get ML encryption master key", e);
Expand Down
199 changes: 193 additions & 6 deletions ml-algorithms/src/test/java/org/opensearch/ml/engine/encryptor/EncryptorImplTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import static org.opensearch.ml.common.CommonValue.CREATE_TIME_FIELD;
Expand All @@ -10,6 +11,7 @@
import static org.opensearch.ml.engine.encryptor.EncryptorImpl.MASTER_KEY_NOT_READY_ERROR;

import java.time.Instant;
import java.util.Map;

import org.junit.Assert;
import org.junit.Before;
Expand All @@ -21,6 +23,7 @@
import org.opensearch.ResourceNotFoundException;
import org.opensearch.Version;
import org.opensearch.action.get.GetResponse;
import org.opensearch.action.index.IndexResponse;
import org.opensearch.client.Client;
import org.opensearch.cluster.ClusterState;
import org.opensearch.cluster.metadata.IndexMetadata;
Expand All @@ -30,6 +33,9 @@
import org.opensearch.common.util.concurrent.ThreadContext;
import org.opensearch.commons.ConfigConstants;
import org.opensearch.core.action.ActionListener;
import org.opensearch.core.index.shard.ShardId;
import org.opensearch.index.engine.VersionConflictEngineException;
import org.opensearch.ml.engine.indices.MLIndicesHandler;
import org.opensearch.threadpool.ThreadPool;

import com.google.common.collect.ImmutableMap;
Expand All @@ -46,6 +52,9 @@ public class EncryptorImplTest {
@Mock
ClusterState clusterState;

@Mock
private MLIndicesHandler mlIndicesHandler;

String masterKey;

@Mock
Expand Down Expand Up @@ -100,14 +109,192 @@ public void setUp() {
}

@Test
public void encrypt() {
Encryptor encryptor = new EncryptorImpl(clusterService, client);
public void encrypt_ExistingMasterKey() {
doAnswer(invocation -> {
ActionListener<Boolean> actionListener = (ActionListener) invocation.getArgument(0);
actionListener.onResponse(true);
return null;
}).when(mlIndicesHandler).initMLConfigIndex(any());
doAnswer(invocation -> {
ActionListener<GetResponse> actionListener = (ActionListener) invocation.getArgument(1);
GetResponse response = mock(GetResponse.class);
when(response.isExists()).thenReturn(true);
when(response.getSourceAsMap()).thenReturn(Map.of(MASTER_KEY, masterKey));
actionListener.onResponse(response);
return null;
}).when(client).get(any(), any());

Encryptor encryptor = new EncryptorImpl(clusterService, client, mlIndicesHandler);
Assert.assertNull(encryptor.getMasterKey());
String encrypted = encryptor.encrypt("test");
Assert.assertNotNull(encrypted);
Assert.assertEquals(masterKey, encryptor.getMasterKey());
}

@Test
public void encrypt_NonExistingMasterKey() {
doAnswer(invocation -> {
ActionListener<Boolean> actionListener = (ActionListener) invocation.getArgument(0);
actionListener.onResponse(true);
return null;
}).when(mlIndicesHandler).initMLConfigIndex(any());
doAnswer(invocation -> {
ActionListener<GetResponse> actionListener = (ActionListener) invocation.getArgument(1);
GetResponse response = mock(GetResponse.class);
when(response.isExists()).thenReturn(false);
actionListener.onResponse(response);
return null;
}).when(client).get(any(), any());
doAnswer(invocation -> {
ActionListener<IndexResponse> actionListener = (ActionListener) invocation.getArgument(1);
IndexResponse response = mock(IndexResponse.class);
actionListener.onResponse(response);
return null;
}).when(client).index(any(), any());

Encryptor encryptor = new EncryptorImpl(clusterService, client, mlIndicesHandler);
Assert.assertNull(encryptor.getMasterKey());
String encrypted = encryptor.encrypt("test");
Assert.assertNotNull(encrypted);
Assert.assertNotEquals(masterKey, encryptor.getMasterKey());
}

@Test
public void encrypt_NonExistingMasterKey_FailedToCreateNewKey() {
exceptionRule.expect(ResourceNotFoundException.class);
exceptionRule.expectMessage(MASTER_KEY_NOT_READY_ERROR);
doAnswer(invocation -> {
ActionListener<Boolean> actionListener = (ActionListener) invocation.getArgument(0);
actionListener.onResponse(true);
return null;
}).when(mlIndicesHandler).initMLConfigIndex(any());
doAnswer(invocation -> {
ActionListener<GetResponse> actionListener = (ActionListener) invocation.getArgument(1);
GetResponse response = mock(GetResponse.class);
when(response.isExists()).thenReturn(false);
actionListener.onResponse(response);
return null;
}).when(client).get(any(), any());
doAnswer(invocation -> {
ActionListener<IndexResponse> actionListener = (ActionListener) invocation.getArgument(1);
actionListener.onFailure(new RuntimeException("random test exception"));
return null;
}).when(client).index(any(), any());

Encryptor encryptor = new EncryptorImpl(clusterService, client, mlIndicesHandler);
Assert.assertNull(encryptor.getMasterKey());
encryptor.encrypt("test");
}

@Test
public void encrypt_NonExistingMasterKey_FailedToCreateNewKey_VersionConflict() {
exceptionRule.expect(ResourceNotFoundException.class);
exceptionRule.expectMessage(MASTER_KEY_NOT_READY_ERROR);
doAnswer(invocation -> {
ActionListener<Boolean> actionListener = (ActionListener) invocation.getArgument(0);
actionListener.onResponse(true);
return null;
}).when(mlIndicesHandler).initMLConfigIndex(any());
doAnswer(invocation -> {
ActionListener<GetResponse> actionListener = (ActionListener) invocation.getArgument(1);
GetResponse response = mock(GetResponse.class);
when(response.isExists()).thenReturn(false);
actionListener.onResponse(response);
return null;
}).doAnswer(invocation -> {
ActionListener<GetResponse> actionListener = (ActionListener) invocation.getArgument(1);
GetResponse response = mock(GetResponse.class);
when(response.isExists()).thenReturn(false);
actionListener.onResponse(response);
return null;
}).when(client).get(any(), any());
doAnswer(invocation -> {
ActionListener<IndexResponse> actionListener = (ActionListener) invocation.getArgument(1);
actionListener
.onFailure(new VersionConflictEngineException(new ShardId(ML_CONFIG_INDEX, "index_uuid", 1), "test_id", "failed"));
return null;
}).when(client).index(any(), any());

Encryptor encryptor = new EncryptorImpl(clusterService, client, mlIndicesHandler);
Assert.assertNull(encryptor.getMasterKey());
encryptor.encrypt("test");
}

@Test
public void encrypt_NonExistingMasterKey_FailedToCreateNewKey_VersionConflict_GetExistingMasterKey() {
doAnswer(invocation -> {
ActionListener<Boolean> actionListener = (ActionListener) invocation.getArgument(0);
actionListener.onResponse(true);
return null;
}).when(mlIndicesHandler).initMLConfigIndex(any());
doAnswer(invocation -> {
ActionListener<GetResponse> actionListener = (ActionListener) invocation.getArgument(1);
GetResponse response = mock(GetResponse.class);
when(response.isExists()).thenReturn(false);
actionListener.onResponse(response);
return null;
}).doAnswer(invocation -> {
ActionListener<GetResponse> actionListener = (ActionListener) invocation.getArgument(1);
GetResponse response = mock(GetResponse.class);
when(response.isExists()).thenReturn(true);
when(response.getSourceAsMap()).thenReturn(Map.of(MASTER_KEY, masterKey));
actionListener.onResponse(response);
return null;
}).when(client).get(any(), any());
doAnswer(invocation -> {
ActionListener<IndexResponse> actionListener = (ActionListener) invocation.getArgument(1);
actionListener
.onFailure(new VersionConflictEngineException(new ShardId(ML_CONFIG_INDEX, "index_uuid", 1), "test_id", "failed"));
return null;
}).when(client).index(any(), any());

Encryptor encryptor = new EncryptorImpl(clusterService, client, mlIndicesHandler);
Assert.assertNull(encryptor.getMasterKey());
String encrypted = encryptor.encrypt("test");
Assert.assertNotNull(encrypted);
Assert.assertEquals(masterKey, encryptor.getMasterKey());
}

@Test
public void encrypt_ThrowExceptionWhenInitMLConfigIndex() {
exceptionRule.expect(RuntimeException.class);
exceptionRule.expectMessage("test exception");
doThrow(new RuntimeException("test exception")).when(mlIndicesHandler).initMLConfigIndex(any());
Encryptor encryptor = new EncryptorImpl(clusterService, client, mlIndicesHandler);
encryptor.encrypt(masterKey);
}

@Test
public void encrypt_FailedToInitMLConfigIndex() {
exceptionRule.expect(ResourceNotFoundException.class);
exceptionRule.expectMessage(MASTER_KEY_NOT_READY_ERROR);
doAnswer(invocation -> {
ActionListener<Boolean> actionListener = (ActionListener) invocation.getArgument(0);
actionListener.onFailure(new RuntimeException("random test exception"));
return null;
}).when(mlIndicesHandler).initMLConfigIndex(any());
Encryptor encryptor = new EncryptorImpl(clusterService, client, mlIndicesHandler);
encryptor.encrypt(masterKey);
}

@Test
public void encrypt_FailedToGetMasterKey() {
exceptionRule.expect(ResourceNotFoundException.class);
exceptionRule.expectMessage(MASTER_KEY_NOT_READY_ERROR);
doAnswer(invocation -> {
ActionListener<Boolean> actionListener = (ActionListener) invocation.getArgument(0);
actionListener.onResponse(true);
return null;
}).when(mlIndicesHandler).initMLConfigIndex(any());
doAnswer(invocation -> {
ActionListener<GetResponse> actionListener = (ActionListener) invocation.getArgument(1);
actionListener.onFailure(new RuntimeException("random test exception"));
return null;
}).when(client).get(any(), any());
Encryptor encryptor = new EncryptorImpl(clusterService, client, mlIndicesHandler);
encryptor.encrypt(masterKey);
}

@Test
public void encrypt_DifferentMasterKey() {
Encryptor encryptor = new EncryptorImpl(masterKey);
Expand All @@ -121,7 +308,7 @@ public void encrypt_DifferentMasterKey() {

@Test
public void decrypt() {
Encryptor encryptor = new EncryptorImpl(clusterService, client);
Encryptor encryptor = new EncryptorImpl(clusterService, client, mlIndicesHandler);
Assert.assertNull(encryptor.getMasterKey());
String encrypted = encryptor.encrypt("test");
String decrypted = encryptor.decrypt(encrypted);
Expand All @@ -142,7 +329,7 @@ public void encrypt_NullMasterKey_NullMasterKey_MasterKeyNotExistInIndex() {
return null;
}).when(client).get(any(), any());

Encryptor encryptor = new EncryptorImpl(clusterService, client);
Encryptor encryptor = new EncryptorImpl(clusterService, client, mlIndicesHandler);
Assert.assertNull(encryptor.getMasterKey());
encryptor.encrypt("test");
}
Expand All @@ -158,7 +345,7 @@ public void decrypt_NullMasterKey_GetMasterKey_Exception() {
return null;
}).when(client).get(any(), any());

Encryptor encryptor = new EncryptorImpl(clusterService, client);
Encryptor encryptor = new EncryptorImpl(clusterService, client, mlIndicesHandler);
Assert.assertNull(encryptor.getMasterKey());
encryptor.decrypt("test");
}
Expand All @@ -177,7 +364,7 @@ public void decrypt_MLConfigIndexNotFound() {
return null;
}).when(client).get(any(), any());

Encryptor encryptor = new EncryptorImpl(clusterService, client);
Encryptor encryptor = new EncryptorImpl(clusterService, client, mlIndicesHandler);
Assert.assertNull(encryptor.getMasterKey());
encryptor.decrypt("test");
}
Expand Down

0 comments on commit 00eb3a1

Please sign in to comment.