Add more UT for remote inference classes (#1077)

(cherry picked from commit daab3b7)

common/build.gradle

@@ -36,11 +36,11 @@ jacocoTestCoverageVerification {
         rule {
             limit {
                 counter = 'LINE'
-                minimum = 0.6 //TODO: add more test to meet the coverage bar 0.9
+                minimum = 0.8 //TODO: add more test to meet the coverage bar 0.9
             }
             limit {
                 counter = 'BRANCH'
-                minimum = 0.5 //TODO: add more test to meet the coverage bar 0.9
+                minimum = 0.7 //TODO: add more test to meet the coverage bar 0.9
             }
         }
     }

common/src/main/java/org/opensearch/ml/common/transport/connector/MLCreateConnectorInput.java

@@ -227,7 +227,7 @@ public void writeTo(StreamOutput output) throws IOException {
         }
         if (!CollectionUtils.isEmpty(backendRoles)) {
             output.writeBoolean(true);
-            output.writeOptionalStringCollection(backendRoles);
+            output.writeStringCollection(backendRoles);
         } else {
             output.writeBoolean(false);
         }

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);
+    }
+}