From f68343678d291b8a5b06cdc1c19dd764fd3eb697 Mon Sep 17 00:00:00 2001
From: Daniel Widdis <widdis@gmail.com>
Date: Wed, 24 Apr 2024 14:53:44 -0700
Subject: [PATCH] Even more test coverage

Signed-off-by: Daniel Widdis <widdis@gmail.com>
---
 .../flowframework/util/EncryptorUtils.java    |  2 +-
 .../util/EncryptorUtilsTests.java             | 30 +++++++++++++++++--
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java b/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java
index 262dc8340..5d9290959 100644
--- 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
index 1815beeb3..eb6ee1544 100644
--- 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() {