add agent framework security it tests

Signed-off-by: Xun Zhang <xunzh@amazon.com>

Author: Zhangxunmt

plugin/src/test/java/org/opensearch/ml/rest/MLCommonsRestTestCase.java

@@ -66,6 +66,8 @@
 import org.opensearch.ml.common.AccessMode;
 import org.opensearch.ml.common.FunctionName;
 import org.opensearch.ml.common.MLTaskState;
+import org.opensearch.ml.common.agent.MLAgent;
+import org.opensearch.ml.common.agent.MLToolSpec;
 import org.opensearch.ml.common.dataset.MLInputDataset;
 import org.opensearch.ml.common.dataset.SearchQueryInputDataset;
 import org.opensearch.ml.common.dataset.TextDocsInputDataSet;
@@ -731,6 +733,32 @@ public String registerModel(String input) throws IOException {
         return parseTaskIdFromResponse(response);
     }
 
+    public void registerMLAgent(RestClient client, String input, Consumer<Map<String, Object>> function) throws IOException {
+        Response response = TestHelper.makeRequest(client, "POST", "/_plugins/_ml/agents/_register", null, input, null);
+        verifyResponse(function, response);
+    }
+
+    public void executeAgent(RestClient client, String agentId, String input, Consumer<Map<String, Object>> function) throws IOException {
+        Response response = TestHelper.makeRequest(client, "POST", "/_plugins/_ml/agents/" + agentId + "/_execute", null, input, null);
+        verifyResponse(function, response);
+    }
+
+    public MLAgent createCatIndexToolMLAgent() {
+        MLToolSpec catIndexTool = MLToolSpec
+            .builder()
+            .type("CatIndexTool")
+            .name("DemoCatIndexTool")
+            .parameters(Map.of("input", "${parameters.question}"))
+            .build();
+        return MLAgent
+            .builder()
+            .name("Test_Agent_For_CatIndex_tool")
+            .type("flow")
+            .description("this is a test agent for the CatIndexTool")
+            .tools(List.of(catIndexTool))
+            .build();
+    }
+
     public void deployModel(RestClient client, MLRegisterModelInput registerModelInput, Consumer<Map<String, Object>> function)
         throws IOException,
         InterruptedException {

plugin/src/test/java/org/opensearch/ml/rest/SecureMLRestIT.java

@@ -26,6 +26,9 @@
 import org.opensearch.ml.common.AccessMode;
 import org.opensearch.ml.common.FunctionName;
 import org.opensearch.ml.common.MLTaskState;
+import org.opensearch.ml.common.agent.MLAgent;
+import org.opensearch.ml.common.dataset.remote.RemoteInferenceInputDataSet;
+import org.opensearch.ml.common.input.execute.agent.AgentMLInput;
 import org.opensearch.ml.common.input.parameter.clustering.KMeansParams;
 import org.opensearch.ml.common.transport.model_group.MLRegisterModelGroupInput;
 import org.opensearch.ml.common.transport.register.MLRegisterModelInput;
@@ -35,6 +38,8 @@
 import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableList;
 
+import static org.opensearch.ml.common.input.execute.agent.AgentMLInput.AgentMLInputBuilder;
+
 public class SecureMLRestIT extends MLCommonsRestTestCase {
     private String irisIndex = "iris_data_secure_ml_it";
 
@@ -248,6 +253,38 @@ public void testDeployModelWithFullAccess() throws IOException, InterruptedExcep
         });
     }
 
+    public void testExecuteAgentWithFullAccess() throws IOException {
+        MLAgent mlAgent = createCatIndexToolMLAgent();
+        registerMLAgent(mlFullAccessClient, TestHelper.toJsonString(mlAgent), registerMLAgentResult -> {
+            assertNotNull(registerMLAgentResult);
+            assertTrue(registerMLAgentResult.containsKey("agent_id"));
+            String agentId = (String) registerMLAgentResult.get("agent_id");
+            try {
+                AgentMLInput agentMLInput = AgentMLInput
+                        .AgentMLInputBuilder()
+                        .agentId(agentId)
+                        .functionName(FunctionName.AGENT)
+                        .inputDataset(RemoteInferenceInputDataSet.builder().parameters(Map.of("question", "How many indices do I have?")).build())
+                        .build();
+
+                executeAgent(mlFullAccessClient, agentId, TestHelper.toJsonString(agentMLInput), mlExecuteTaskResponse -> {
+                    assertNotNull(mlExecuteTaskResponse);
+                    assertTrue(mlExecuteTaskResponse.containsKey("output"));
+                });
+            } catch (IOException e) {
+                assertNull(e);
+            }
+        });
+    }
+
+    public void testRegisterAgentWithFullAccess() throws IOException {
+        MLAgent mlAgent = createCatIndexToolMLAgent();
+        registerMLAgent(mlFullAccessClient, TestHelper.toJsonString(mlAgent), registerMLAgentResult -> {
+            assertNotNull(registerMLAgentResult);
+            assertTrue(registerMLAgentResult.containsKey("agent_id"));
+        });
+    }
+
     public void testTrainWithReadOnlyMLAccess() throws IOException {
         exceptionRule.expect(ResponseException.class);
         exceptionRule.expectMessage("no permissions for [cluster:admin/opensearch/ml/train]");