Even more test coverage

Signed-off-by: Daniel Widdis <[email protected]>
opensearch-project · Apr 26, 2024 · 3818ca6 · 3818ca6
1 parent 686b021
commit 3818ca6
Showing 1 changed file with 26 additions and 0 deletions.
diff --git a/src/test/java/org/opensearch/flowframework/util/EncryptorUtilsTests.java b/src/test/java/org/opensearch/flowframework/util/EncryptorUtilsTests.java
@@ -11,6 +11,8 @@
 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;
@@ -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() {