Add tests for rejected master key

Signed-off-by: Ashish Agrawal <[email protected]>
opensearch-project · Aug 28, 2024 · 935b606 · 935b606
1 parent 55e0cd1
commit 935b606
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 0 deletions.
diff --git a/client/src/test/java/org/opensearch/ml/client/MachineLearningNodeClientTest.java b/client/src/test/java/org/opensearch/ml/client/MachineLearningNodeClientTest.java
@@ -14,6 +14,7 @@
 import static org.mockito.ArgumentMatchers.isA;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.verify;
+import static org.opensearch.ml.common.CommonValue.MASTER_KEY;
 import static org.opensearch.ml.common.input.Constants.ACTION;
 import static org.opensearch.ml.common.input.Constants.ALGORITHM;
 import static org.opensearch.ml.common.input.Constants.KMEANS;
@@ -40,6 +41,7 @@
 import org.mockito.InjectMocks;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.opensearch.OpenSearchStatusException;
 import org.opensearch.action.delete.DeleteResponse;
 import org.opensearch.action.search.SearchRequest;
 import org.opensearch.action.search.SearchResponse;
@@ -51,6 +53,7 @@
 import org.opensearch.core.common.bytes.BytesReference;
 import org.opensearch.core.index.Index;
 import org.opensearch.core.index.shard.ShardId;
+import org.opensearch.core.rest.RestStatus;
 import org.opensearch.core.xcontent.ToXContent;
 import org.opensearch.core.xcontent.ToXContentObject;
 import org.opensearch.core.xcontent.XContentBuilder;
@@ -979,6 +982,23 @@ public void getConfig() {
         assertEquals("type", argumentCaptor.getValue().getType());
     }
 
+    @Test
+    public void getConfigRejectedMasterKey() {
+        doAnswer(invocation -> {
+            ActionListener<MLConfigGetResponse> actionListener = invocation.getArgument(2);
+            actionListener.onFailure(new OpenSearchStatusException("You are not allowed to access this config doc", RestStatus.FORBIDDEN));
+            return null;
+        }).when(client).execute(eq(MLConfigGetAction.INSTANCE), any(), any());
+
+        ArgumentCaptor<OpenSearchStatusException> argumentCaptor = ArgumentCaptor.forClass(OpenSearchStatusException.class);
+        machineLearningNodeClient.getConfig(MASTER_KEY, getMlConfigListener);
+
+        verify(client).execute(eq(MLConfigGetAction.INSTANCE), isA(MLConfigGetRequest.class), any());
+        verify(getMlConfigListener).onFailure(argumentCaptor.capture());
+        assertEquals(RestStatus.FORBIDDEN, argumentCaptor.getValue().status());
+        assertEquals("You are not allowed to access this config doc", argumentCaptor.getValue().getLocalizedMessage());
+    }
+
     private SearchResponse createSearchResponse(ToXContentObject o) throws IOException {
         XContentBuilder content = o.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS);
 

diff --git a/plugin/src/test/java/org/opensearch/ml/action/config/GetConfigTransportActionTests.java b/plugin/src/test/java/org/opensearch/ml/action/config/GetConfigTransportActionTests.java
@@ -5,12 +5,14 @@
 
 package org.opensearch.ml.action.config;
 
+import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
+import static org.opensearch.ml.common.CommonValue.MASTER_KEY;
 
 import java.io.IOException;
 import java.time.Instant;
@@ -22,6 +24,7 @@
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.opensearch.OpenSearchStatusException;
 import org.opensearch.action.get.GetResponse;
 import org.opensearch.action.support.ActionFilters;
 import org.opensearch.client.Client;
@@ -30,6 +33,7 @@
 import org.opensearch.common.xcontent.XContentFactory;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.common.bytes.BytesReference;
+import org.opensearch.core.rest.RestStatus;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.core.xcontent.ToXContent;
 import org.opensearch.core.xcontent.XContentBuilder;
@@ -168,4 +172,27 @@ public GetResponse prepareMLConfig(String configID) throws IOException {
         GetResponse getResponse = new GetResponse(getResult);
         return getResponse;
     }
+
+    @Test
+    public void testDoExecute_Rejected_MASTER_KEY() throws IOException {
+        String configID = MASTER_KEY;
+        GetResponse getResponse = prepareMLConfig(configID);
+        ActionListener<MLConfigGetResponse> actionListener = mock(ActionListener.class);
+        MLConfigGetRequest request = new MLConfigGetRequest(configID);
+        Task task = mock(Task.class);
+
+        doAnswer(invocation -> {
+            ActionListener<GetResponse> listener = invocation.getArgument(1);
+            listener.onResponse(getResponse);
+            return null;
+        }).when(client).get(any(), any());
+
+        ArgumentCaptor<OpenSearchStatusException> argumentCaptor = ArgumentCaptor.forClass(OpenSearchStatusException.class);
+
+        getConfigTransportAction.doExecute(task, request, actionListener);
+        verify(actionListener).onFailure(argumentCaptor.capture());
+        assertEquals(RestStatus.FORBIDDEN, argumentCaptor.getValue().status());
+        assertEquals("You are not allowed to access this config doc", argumentCaptor.getValue().getLocalizedMessage());
+
+    }
 }