Skip to content

Commit

Permalink
Add more UT for remote inference classes (opensearch-project#1077)
Browse files Browse the repository at this point in the history
  • Loading branch information
b4sjoo authored and zane-neo committed Sep 1, 2023
1 parent bfc6694 commit dacd510
Show file tree
Hide file tree
Showing 13 changed files with 742 additions and 86 deletions.
4 changes: 2 additions & 2 deletions common/build.gradle
Original file line number Diff line number Diff line change
Expand Up @@ -40,11 +40,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
}
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down
Original file line number Diff line number Diff line change
@@ -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());
}
}
Original file line number Diff line number Diff line change
@@ -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);
}
}

Original file line number Diff line number Diff line change
@@ -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);
}
}
Loading

0 comments on commit dacd510

Please sign in to comment.