Even more test coverage

Signed-off-by: Daniel Widdis <[email protected]>
opensearch-project · Apr 27, 2024 · f683436 · f683436
1 parent 5d91205
commit f683436
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 3 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java b/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java
@@ -17,8 +17,8 @@
 import org.opensearch.client.Client;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.util.concurrent.ThreadContext;
-import org.opensearch.commons.authuser.User;
 import org.opensearch.common.xcontent.XContentFactory;
+import org.opensearch.commons.authuser.User;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.rest.RestStatus;
 import org.opensearch.core.xcontent.NamedXContentRegistry;

diff --git a/src/test/java/org/opensearch/flowframework/util/EncryptorUtilsTests.java b/src/test/java/org/opensearch/flowframework/util/EncryptorUtilsTests.java
@@ -11,14 +11,16 @@
 import org.opensearch.Version;
 import org.opensearch.action.get.GetRequest;
 import org.opensearch.action.get.GetResponse;
+import org.opensearch.action.index.IndexRequest;
+import org.opensearch.action.index.IndexResponse;
 import org.opensearch.client.Client;
 import org.opensearch.cluster.ClusterState;
 import org.opensearch.cluster.metadata.Metadata;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.util.concurrent.ThreadContext;
-import org.opensearch.commons.authuser.User;
 import org.opensearch.common.xcontent.XContentFactory;
+import org.opensearch.commons.authuser.User;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.common.bytes.BytesReference;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
@@ -33,9 +35,9 @@
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.threadpool.ThreadPool;
 
-import java.util.ArrayList;
 import java.io.IOException;
 import java.time.Instant;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -140,6 +142,8 @@ public void testInitializeMasterKeySuccess() throws IOException {
         encryptorUtils.setMasterKey(null);
 
         String masterKey = encryptorUtils.generateMasterKey();
+
+        // Index exists case
         BytesReference bytesRef;
         try (XContentBuilder builder = XContentFactory.jsonBuilder()) {
             Config config = new Config(masterKey, Instant.now());
@@ -162,11 +166,33 @@ public void testInitializeMasterKeySuccess() throws IOException {
         encryptorUtils.initializeMasterKey(listener);
         assertEquals(masterKey, encryptorUtils.getMasterKey());
 
+        // Test ifAbsent version
         encryptorUtils.setMasterKey(null);
         assertNull(encryptorUtils.getMasterKey());
 
         encryptorUtils.initializeMasterKeyIfAbsent();
         assertEquals(masterKey, encryptorUtils.getMasterKey());
+
+        // No index exists case
+        doAnswer(invocation -> {
+            ActionListener<GetResponse> getRequestActionListener = invocation.getArgument(1);
+            GetResponse getResponse = mock(GetResponse.class);
+            when(getResponse.isExists()).thenReturn(false);
+            getRequestActionListener.onResponse(getResponse);
+            return null;
+        }).when(client).get(any(GetRequest.class), any());
+        doAnswer(invocation -> {
+            ActionListener<IndexResponse> indexRequestActionListener = invocation.getArgument(1);
+            IndexResponse indexResponse = mock(IndexResponse.class);
+            indexRequestActionListener.onResponse(indexResponse);
+            return null;
+        }).when(client).index(any(IndexRequest.class), any());
+
+        listener = ActionListener.wrap(b -> {}, e -> {});
+        encryptorUtils.initializeMasterKey(listener);
+        // This will generate a new master key 32 bytes -> base64 encoded
+        assertNotEquals(masterKey, encryptorUtils.getMasterKey());
+        assertEquals(masterKey.length(), encryptorUtils.getMasterKey().length());
     }
 
     public void testInitializeMasterKeyFailure() {