Add more UT for remote inference classes

common/src/test/java/org/opensearch/ml/common/transport/connector/MLConnectorDeleteRequestTests.java

@@ -0,0 +1,100 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.transport.connector;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.opensearch.action.ActionRequest;
+import org.opensearch.action.ActionRequestValidationException;
+import org.opensearch.common.io.stream.BytesStreamOutput;
+import org.opensearch.common.io.stream.StreamOutput;
+
+import java.io.IOException;
+import java.io.UncheckedIOException;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
+
+public class MLConnectorDeleteRequestTests {
+    private String connectorId;
+
+    @Before
+    public void setUp() {
+        connectorId = "test-connector-id";
+    }
+
+    @Test
+    public void writeTo_Success() throws IOException {
+        MLConnectorDeleteRequest mlConnectorDeleteRequest = MLConnectorDeleteRequest.builder()
+                .connectorId(connectorId).build();
+        BytesStreamOutput bytesStreamOutput = new BytesStreamOutput();
+        mlConnectorDeleteRequest.writeTo(bytesStreamOutput);
+        MLConnectorDeleteRequest parsedConnector = new MLConnectorDeleteRequest(bytesStreamOutput.bytes().streamInput());
+        assertEquals(parsedConnector.getConnectorId(), connectorId);
+    }
+
+    @Test
+    public void valid_Exception_NullConnectorId() {
+        MLConnectorDeleteRequest mlConnectorDeleteRequest = MLConnectorDeleteRequest.builder().build();
+        ActionRequestValidationException exception = mlConnectorDeleteRequest.validate();
+        assertEquals("Validation Failed: 1: ML connector id can't be null;", exception.getMessage());
+    }
+
+    @Test
+    public void validate_Success() {
+        MLConnectorDeleteRequest mlConnectorDeleteRequest = MLConnectorDeleteRequest.builder()
+                .connectorId(connectorId).build();
+        ActionRequestValidationException actionRequestValidationException = mlConnectorDeleteRequest.validate();
+        assertNull(actionRequestValidationException);
+    }
+
+    @Test
+    public void fromActionRequest_Success() {
+        MLConnectorDeleteRequest mlConnectorDeleteRequest = MLConnectorDeleteRequest.builder()
+                .connectorId(connectorId).build();
+        ActionRequest actionRequest = new ActionRequest() {
+            @Override
+            public ActionRequestValidationException validate() {
+                return null;
+            }
+
+            @Override
+            public void writeTo(StreamOutput out) throws IOException {
+                mlConnectorDeleteRequest.writeTo(out);
+            }
+        };
+        MLConnectorDeleteRequest parsedConnector = MLConnectorDeleteRequest.fromActionRequest(actionRequest);
+        assertNotSame(parsedConnector, mlConnectorDeleteRequest);
+        assertEquals(parsedConnector.getConnectorId(), connectorId);
+    }
+
+    @Test(expected = UncheckedIOException.class)
+    public void fromActionRequest_IOException() {
+        ActionRequest actionRequest = new ActionRequest() {
+            @Override
+            public ActionRequestValidationException validate() {
+                return null;
+            }
+
+            @Override
+            public void writeTo(StreamOutput out) throws IOException {
+                throw new IOException();
+            }
+        };
+        MLConnectorDeleteRequest.fromActionRequest(actionRequest);
+    }
+
+    @Test
+    public void fromActionRequestWithConnectorDeleteRequest_Success() {
+        MLConnectorDeleteRequest mlConnectorDeleteRequest = MLConnectorDeleteRequest.builder()
+                .connectorId(connectorId).build();
+        MLConnectorDeleteRequest mlConnectorDeleteRequestFromActionRequest = MLConnectorDeleteRequest.fromActionRequest(mlConnectorDeleteRequest);
+        assertSame(mlConnectorDeleteRequest, mlConnectorDeleteRequestFromActionRequest);
+        assertEquals(mlConnectorDeleteRequest.getConnectorId(), mlConnectorDeleteRequestFromActionRequest.getConnectorId());
+    }
+}

common/src/test/java/org/opensearch/ml/common/transport/connector/MLConnectorGetRequestTests.java

@@ -0,0 +1,98 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+
+package org.opensearch.ml.common.transport.connector;
+
+import java.io.IOException;
+import java.io.UncheckedIOException;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.opensearch.action.ActionRequest;
+import org.opensearch.action.ActionRequestValidationException;
+import org.opensearch.common.io.stream.BytesStreamOutput;
+import org.opensearch.common.io.stream.StreamOutput;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
+
+public class MLConnectorGetRequestTests {
+    private String connectorId;
+
+    @Before
+    public void setUp() {
+        connectorId = "test-connector-id";
+    }
+
+    @Test
+    public void writeTo_Success() throws IOException {
+        MLConnectorGetRequest mlConnectorGetRequest = MLConnectorGetRequest.builder().connectorId(connectorId).build();
+        BytesStreamOutput bytesStreamOutput = new BytesStreamOutput();
+        mlConnectorGetRequest.writeTo(bytesStreamOutput);
+        MLConnectorGetRequest parsedConnector = new MLConnectorGetRequest(bytesStreamOutput.bytes().streamInput());
+        assertEquals(connectorId, parsedConnector.getConnectorId());
+    }
+
+    @Test
+    public void fromActionRequest_Success() {
+        MLConnectorGetRequest mlConnectorGetRequest = MLConnectorGetRequest.builder().connectorId(connectorId).build();
+        ActionRequest actionRequest = new ActionRequest() {
+            @Override
+            public ActionRequestValidationException validate() {
+                return null;
+            }
+
+            @Override
+            public void writeTo(StreamOutput out) throws IOException {
+                mlConnectorGetRequest.writeTo(out);
+            }
+        };
+        MLConnectorGetRequest mlConnectorGetRequestFromActionRequest = MLConnectorGetRequest.fromActionRequest(actionRequest);
+        assertNotSame(mlConnectorGetRequest, mlConnectorGetRequestFromActionRequest);
+        assertEquals(mlConnectorGetRequest.getConnectorId(), mlConnectorGetRequestFromActionRequest.getConnectorId());
+    }
+
+    @Test(expected = UncheckedIOException.class)
+    public void fromActionRequest_IOException() {
+        ActionRequest actionRequest = new ActionRequest() {
+            @Override
+            public ActionRequestValidationException validate() {
+                return null;
+            }
+
+            @Override
+            public void writeTo(StreamOutput out) throws IOException {
+                throw new IOException();
+            }
+        };
+        MLConnectorGetRequest.fromActionRequest(actionRequest);
+    }
+
+    @Test
+    public void fromActionRequestWithMLConnectorGetRequest_Success() {
+        MLConnectorGetRequest mlConnectorGetRequest = MLConnectorGetRequest.builder().connectorId(connectorId).build();
+        MLConnectorGetRequest mlConnectorGetRequestFromActionRequest = MLConnectorGetRequest.fromActionRequest(mlConnectorGetRequest);
+        assertSame(mlConnectorGetRequest, mlConnectorGetRequestFromActionRequest);
+        assertEquals(mlConnectorGetRequest.getConnectorId(), mlConnectorGetRequestFromActionRequest.getConnectorId());
+    }
+
+    @Test
+    public void validate_Exception_NullConnctorId() {
+        MLConnectorGetRequest mlConnectorGetRequest = MLConnectorGetRequest.builder().build();
+        ActionRequestValidationException actionRequestValidationException = mlConnectorGetRequest.validate();
+        assertEquals("Validation Failed: 1: ML connector id can't be null;", actionRequestValidationException.getMessage());
+    }
+
+    @Test
+    public void validate_Success() {
+        MLConnectorGetRequest mlConnectorGetRequest = MLConnectorGetRequest.builder().connectorId(connectorId).build();
+        ActionRequestValidationException actionRequestValidationException = mlConnectorGetRequest.validate();
+        assertNull(actionRequestValidationException);
+    }
+}
+

common/src/test/java/org/opensearch/ml/common/transport/connector/MLConnectorGetResponseTests.java

@@ -0,0 +1,107 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.transport.connector;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.opensearch.action.ActionResponse;
+import org.opensearch.common.Strings;
+import org.opensearch.common.io.stream.BytesStreamOutput;
+import org.opensearch.common.io.stream.StreamOutput;
+import org.opensearch.common.xcontent.XContentFactory;
+import org.opensearch.common.xcontent.XContentType;
+import org.opensearch.core.xcontent.ToXContent;
+import org.opensearch.core.xcontent.XContentBuilder;
+import org.opensearch.ml.common.connector.Connector;
+import org.opensearch.ml.common.connector.HttpConnectorTest;
+
+import java.io.IOException;
+import java.io.UncheckedIOException;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertSame;
+
+public class MLConnectorGetResponseTests {
+    Connector connector;
+
+    @Before
+    public void setUp() {
+        connector = HttpConnectorTest.createHttpConnector();
+    }
+
+    @Test
+    public void writeTo_Success() throws IOException {
+        BytesStreamOutput bytesStreamOutput = new BytesStreamOutput();
+        MLConnectorGetResponse response = MLConnectorGetResponse.builder().mlConnector(connector).build();
+        response.writeTo(bytesStreamOutput);
+        MLConnectorGetResponse parsedResponse = new MLConnectorGetResponse(bytesStreamOutput.bytes().streamInput());
+        assertNotEquals(response, parsedResponse);
+        assertNotSame(response.mlConnector, parsedResponse.mlConnector);
+        assertEquals(response.mlConnector, parsedResponse.mlConnector);
+        assertEquals(response.mlConnector.getName(), parsedResponse.mlConnector.getName());
+        assertEquals(response.mlConnector.getAccess(), parsedResponse.mlConnector.getAccess());
+        assertEquals(response.mlConnector.getProtocol(), parsedResponse.mlConnector.getProtocol());
+        assertEquals(response.mlConnector.getDecryptedHeaders(), parsedResponse.mlConnector.getDecryptedHeaders());
+        assertEquals(response.mlConnector.getBackendRoles(), parsedResponse.mlConnector.getBackendRoles());
+        assertEquals(response.mlConnector.getActions(), parsedResponse.mlConnector.getActions());
+        assertEquals(response.mlConnector.getParameters(), parsedResponse.mlConnector.getParameters());
+    }
+
+    @Test
+    public void toXContentTest() throws IOException {
+        MLConnectorGetResponse mlConnectorGetResponse = MLConnectorGetResponse.builder().mlConnector(connector).build();
+        XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON);
+        mlConnectorGetResponse.toXContent(builder, ToXContent.EMPTY_PARAMS);
+        assertNotNull(builder);
+        String jsonStr = Strings.toString(builder);
+        assertEquals("{\"name\":\"test_connector_name\"," +
+                "\"version\":\"1\",\"description\":\"this is a test connector\",\"protocol\":\"http\"," +
+                "\"parameters\":{\"input\":\"test input value\"},\"credential\":{\"key\":\"test_key_value\"}," +
+                "\"actions\":[{\"action_type\":\"PREDICT\",\"method\":\"POST\",\"url\":\"https://test.com\"," +
+                "\"headers\":{\"api_key\":\"${credential.key}\"}," +
+                "\"request_body\":\"{\\\"input\\\": \\\"${parameters.input}\\\"}\"," +
+                "\"pre_process_function\":\"connector.pre_process.openai.embedding\"," +
+                "\"post_process_function\":\"connector.post_process.openai.embedding\"}]," +
+                "\"backend_roles\":[\"role1\",\"role2\"]," +
+                "\"access\":\"public\"}", jsonStr);
+    }
+
+    @Test
+    public void fromActionResponseWithMLConnectorGetResponse_Success() {
+        MLConnectorGetResponse mlConnectorGetResponse = MLConnectorGetResponse.builder().mlConnector(connector).build();
+        MLConnectorGetResponse mlConnectorGetResponseFromActionResponse = MLConnectorGetResponse.fromActionResponse(mlConnectorGetResponse);
+        assertSame(mlConnectorGetResponse, mlConnectorGetResponseFromActionResponse);
+        assertEquals(mlConnectorGetResponse.mlConnector, mlConnectorGetResponseFromActionResponse.mlConnector);
+    }
+
+    @Test
+    public void fromActionResponse_Success() {
+        MLConnectorGetResponse mlConnectorGetResponse = MLConnectorGetResponse.builder().mlConnector(connector).build();
+        ActionResponse actionResponse = new ActionResponse() {
+            @Override
+            public void writeTo(StreamOutput out) throws IOException {
+                mlConnectorGetResponse.writeTo(out);
+            }
+        };
+        MLConnectorGetResponse mlConnectorGetResponseFromActionResponse = MLConnectorGetResponse.fromActionResponse(actionResponse);
+        assertNotSame(mlConnectorGetResponse, mlConnectorGetResponseFromActionResponse);
+        assertEquals(mlConnectorGetResponse.mlConnector, mlConnectorGetResponseFromActionResponse.mlConnector);
+    }
+
+    @Test(expected = UncheckedIOException.class)
+    public void fromActionResponse_IOException() {
+        ActionResponse actionResponse = new ActionResponse() {
+            @Override
+            public void writeTo(StreamOutput out) throws IOException {
+                throw new IOException();
+            }
+        };
+        MLConnectorGetResponse.fromActionResponse(actionResponse);
+    }
+}

common/src/test/java/org/opensearch/ml/common/transport/connector/MLCreateConnectorInputTests.java

@@ -0,0 +1,70 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.transport.connector;
+
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.opensearch.action.ActionRequest;
+import org.opensearch.action.ActionRequestValidationException;
+import org.opensearch.common.io.stream.BytesStreamOutput;
+import org.opensearch.common.io.stream.StreamOutput;
+import org.opensearch.ml.common.AccessMode;
+import org.opensearch.ml.common.connector.ConnectorAction;
+import org.opensearch.ml.common.connector.MLPostProcessFunction;
+import org.opensearch.ml.common.connector.MLPreProcessFunction;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
+
+public class MLCreateConnectorInputTests {
+    private MLCreateConnectorInput mlCreateConnectorInput;
+
+    @Before
+    public void setUp(){
+        ConnectorAction.ActionType actionType = ConnectorAction.ActionType.PREDICT;
+        String method = "POST";
+        String url = "https://test.com";
+        Map<String, String> headers = new HashMap<>();
+        headers.put("api_key", "${credential.key}");
+        String mlCreateConnectorRequestBody = "{\"input\": \"${parameters.input}\"}";
+        String preProcessFunction = MLPreProcessFunction.TEXT_DOCS_TO_OPENAI_EMBEDDING_INPUT;
+        String postProcessFunction = MLPostProcessFunction.OPENAI_EMBEDDING;
+        ConnectorAction action = new ConnectorAction(actionType, method, url, headers, mlCreateConnectorRequestBody, preProcessFunction, postProcessFunction);
+
+        // java.lang.IllegalStateException: unexpected byte [0x6f] will be thrown if we specify backendRoles field.
+        mlCreateConnectorInput = MLCreateConnectorInput.builder()
+                .name("test_connector_name")
+                .description("this is a test connector")
+                .version("1")
+                .protocol("http")
+                .parameters(Map.of("input", "test input value"))
+                .credential(Map.of("key", "test_key_value"))
+                .actions(List.of(action))
+                .access(AccessMode.PUBLIC)
+                .addAllBackendRoles(true)
+                .build();
+    }
+
+    @Test
+    // MLCreateConnectorInput check its parameters when created, so exception is not thrown here
+    public void validate_Exception_NullMLModelName() {
+        mlCreateConnectorInput.setName(null);
+        MLCreateConnectorRequest mlCreateConnectorRequest = MLCreateConnectorRequest.builder()
+                .mlCreateConnectorInput(mlCreateConnectorInput)
+                .build();
+
+        assertNull(mlCreateConnectorRequest.validate());
+        assertNull(mlCreateConnectorRequest.getMlCreateConnectorInput().getName());
+    }
+}