fix npe issue (#2145) (#2194)

* fix npe issue Signed-off-by: Jing Zhang <[email protected]> * add UT and change error message Signed-off-by: Jing Zhang <[email protected]> * add more UT Signed-off-by: Jing Zhang <[email protected]> --------- Signed-off-by: Jing Zhang <[email protected]> (cherry picked from commit 0b07a99) Co-authored-by: Jing Zhang <[email protected]>
opensearch-project · Mar 12, 2024 · 01bf7c0 · 01bf7c0
1 parent 63c72d0
commit 01bf7c0
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLAgentExecutor.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLAgentExecutor.java
@@ -98,8 +98,8 @@ public void execute(Input input, ActionListener<Output> listener) {
         AgentMLInput agentMLInput = (AgentMLInput) input;
         String agentId = agentMLInput.getAgentId();
         RemoteInferenceInputDataSet inputDataSet = (RemoteInferenceInputDataSet) agentMLInput.getInputDataset();
-        if (inputDataSet.getParameters() == null) {
-            throw new IllegalArgumentException("wrong input");
+        if (inputDataSet == null || inputDataSet.getParameters() == null) {
+            throw new IllegalArgumentException("Agent input data can not be empty.");
         }
 
         List<ModelTensors> outputs = new ArrayList<>();

diff --git a/...gorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLAgentExecutorTest.java b/...gorithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLAgentExecutorTest.java
@@ -183,6 +183,19 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         mlAgentExecutor.execute(input, agentActionListener);
     }
 
+    @Test(expected = IllegalArgumentException.class)
+    public void test_NonInputData_ThrowsException() {
+        AgentMLInput agentMLInput = new AgentMLInput("test", FunctionName.AGENT, null);
+        mlAgentExecutor.execute(agentMLInput, agentActionListener);
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void test_NonInputParas_ThrowsException() {
+        RemoteInferenceInputDataSet inputDataSet = RemoteInferenceInputDataSet.builder().parameters(null).build();
+        AgentMLInput agentMLInput = new AgentMLInput("test", FunctionName.AGENT, inputDataSet);
+        mlAgentExecutor.execute(agentMLInput, agentActionListener);
+    }
+
     @Test
     public void test_HappyCase_ReturnsResult() {
         ModelTensor modelTensor = ModelTensor.builder().name("response").dataAsMap(ImmutableMap.of("test_key", "test_value")).build();