add more unit test

Signed-off-by: Yaliang Wu <[email protected]>

Author: ylwu-amzn

common/src/main/java/org/opensearch/ml/common/output/model/ModelTensor.java

@@ -296,7 +296,7 @@ public String toString() {
         try {
             return this.toXContent(JsonXContent.contentBuilder(), null).toString();
         } catch (IOException e) {
-            throw new RuntimeException(e);
+            throw new IllegalArgumentException("Can't convert ModelTensor to string", e);
         }
     }
 

common/src/main/java/org/opensearch/ml/common/output/model/ModelTensors.java

@@ -178,7 +178,7 @@ public String toString() {
         try {
             return this.toXContent(JsonXContent.contentBuilder(), null).toString();
         } catch (IOException e) {
-            throw new RuntimeException(e);
+            throw new IllegalArgumentException("Can't convert ModelTensor to string", e);
         }
     }
 }

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/ConnectorUtils.java

@@ -167,6 +167,9 @@ private static MLInput escapeMLInput(MLInput mlInput) {
     }
 
     public static void escapeRemoteInferenceInputData(RemoteInferenceInputDataSet inputData) {
+        if (inputData.getParameters() == null) {
+            return;
+        }
         Map<String, String> newParameters = new HashMap<>();
         String noEscapeParams = inputData.getParameters().get(NO_ESCAPE_PARAMS);
         Set<String> noEscapParamSet = new HashSet<>();

ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/remote/ConnectorUtilsTest.java

@@ -423,17 +423,69 @@ public void testEscapeRemoteInferenceInputData_WithNoEscapeParams() {
         params.put("key1", inputKey1);
         params.put("key2", "test value");
         params.put("key3", inputKey3);
-        params.put("NO_ESCAPE_PARAMS", "key1,key3");
+        params.put("no_escape_params", "key1,key3");
 
         RemoteInferenceInputDataSet inputData = RemoteInferenceInputDataSet.builder().parameters(params).build();
 
         ConnectorUtils.escapeRemoteInferenceInputData(inputData);
 
-        String expectedKey1 = "hello \\\"world\\\"";
-        String expectedKey3 = "special \\\"chars\\\"";
-        assertEquals(expectedKey1, inputData.getParameters().get("key1"));
+        assertEquals(inputKey1, inputData.getParameters().get("key1"));
         assertEquals("test value", inputData.getParameters().get("key2"));
-        assertEquals(expectedKey3, inputData.getParameters().get("key3"));
+        assertEquals(inputKey3, inputData.getParameters().get("key3"));
+    }
+
+    @Test
+    public void testEscapeRemoteInferenceInputData_NullValue() {
+        Map<String, String> params = new HashMap<>();
+        params.put("key1", null);
+        RemoteInferenceInputDataSet inputData = RemoteInferenceInputDataSet.builder().parameters(params).build();
+
+        ConnectorUtils.escapeRemoteInferenceInputData(inputData);
+
+        assertNull(inputData.getParameters().get("key1"));
+    }
+
+    @Test
+    public void testEscapeRemoteInferenceInputData_JsonValue() {
+        Map<String, String> params = new HashMap<>();
+        params.put("key1", "{\"test\": \"value\"}");
+        RemoteInferenceInputDataSet inputData = RemoteInferenceInputDataSet.builder().parameters(params).build();
+
+        ConnectorUtils.escapeRemoteInferenceInputData(inputData);
+
+        assertEquals("{\"test\": \"value\"}", inputData.getParameters().get("key1"));
+    }
+
+    @Test
+    public void testEscapeRemoteInferenceInputData_EscapeValue() {
+        Map<String, String> params = new HashMap<>();
+        params.put("key1", "test\"value");
+        RemoteInferenceInputDataSet inputData = RemoteInferenceInputDataSet.builder().parameters(params).build();
+
+        ConnectorUtils.escapeRemoteInferenceInputData(inputData);
+
+        assertEquals("test\\\"value", inputData.getParameters().get("key1"));
+    }
+
+    @Test
+    public void testEscapeRemoteInferenceInputData_NoEscapeParam() {
+        Map<String, String> params = new HashMap<>();
+        params.put("key1", "test\"value");
+        params.put("no_escape_params", "key1");
+        RemoteInferenceInputDataSet inputData = RemoteInferenceInputDataSet.builder().parameters(params).build();
+
+        ConnectorUtils.escapeRemoteInferenceInputData(inputData);
+
+        assertEquals("test\"value", inputData.getParameters().get("key1"));
+    }
+
+    @Test
+    public void testEscapeRemoteInferenceInputData_NullParameters() {
+        RemoteInferenceInputDataSet inputData = RemoteInferenceInputDataSet.builder().parameters(null).build();
+
+        ConnectorUtils.escapeRemoteInferenceInputData(inputData);
+
+        assertNull(inputData.getParameters());
     }
 
     @Test
@@ -538,4 +590,58 @@ public void processOutput_WithResponseFilterOnly() throws IOException {
         assertEquals(1, tensors.getMlModelTensors().size());
         assertEquals("response", tensors.getMlModelTensors().get(0).getName());
     }
+
+    @Test
+    public void processOutput_ScriptReturnModelTensor_WithJsonResponse() throws IOException {
+        String postprocessResult = "{\"name\":\"test\",\"data\":[1,2,3]}";
+        when(scriptService.compile(any(), any())).then(invocation -> new TestTemplateService.MockTemplateScript.Factory(postprocessResult));
+
+        ConnectorAction predictAction = ConnectorAction
+            .builder()
+            .actionType(PREDICT)
+            .method("POST")
+            .url("http://test.com/mock")
+            .requestBody("{\"input\": \"${parameters.input}\"}")
+            .postProcessFunction("custom_script")
+            .build();
+        Connector connector = HttpConnector
+            .builder()
+            .name("test connector")
+            .version("1")
+            .protocol("http")
+            .actions(Arrays.asList(predictAction))
+            .build();
+        String modelResponse = "{\"result\":\"test\"}";
+
+        ModelTensors tensors = ConnectorUtils
+            .processOutput(PREDICT.name(), modelResponse, connector, scriptService, ImmutableMap.of(), null);
+
+        assertEquals(1, tensors.getMlModelTensors().size());
+    }
+
+    @Test
+    public void processOutput_WithProcessorChain_StringOutput() throws IOException {
+        ConnectorAction predictAction = ConnectorAction
+            .builder()
+            .actionType(PREDICT)
+            .method("POST")
+            .url("http://test.com/mock")
+            .requestBody("{\"input\": \"${parameters.input}\"}")
+            .build();
+        Map<String, String> parameters = new HashMap<>();
+        parameters.put("processor_configs", "[{\"type\":\"test_processor\"}]");
+        Connector connector = HttpConnector
+            .builder()
+            .name("test connector")
+            .version("1")
+            .protocol("http")
+            .actions(Arrays.asList(predictAction))
+            .build();
+        String modelResponse = "{\"result\":\"test response\"}";
+
+        ModelTensors tensors = ConnectorUtils.processOutput(PREDICT.name(), modelResponse, connector, scriptService, parameters, null);
+
+        assertEquals(1, tensors.getMlModelTensors().size());
+        assertEquals("response", tensors.getMlModelTensors().get(0).getName());
+    }
 }