Skip to content

Commit

Permalink
more ut test coverage
Browse files Browse the repository at this point in the history
Signed-off-by: Xun Zhang <[email protected]>
  • Loading branch information
Zhangxunmt authored and rbhavna committed Nov 16, 2023
1 parent 82090ad commit 9c70503
Show file tree
Hide file tree
Showing 6 changed files with 529 additions and 25 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,128 @@
/*
* 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.mockito.Mock;
import org.mockito.MockitoAnnotations;
import org.opensearch.action.ActionRequest;
import org.opensearch.action.ActionRequestValidationException;
import org.opensearch.common.io.stream.BytesStreamOutput;
import org.opensearch.core.common.io.stream.StreamOutput;
import org.opensearch.core.xcontent.XContentParser;
import org.opensearch.rest.RestRequest;

import java.io.IOException;
import java.io.UncheckedIOException;
import java.util.Map;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotSame;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertSame;
import static org.mockito.Mockito.when;

public class MLUpdateConnectorRequestTests {
private String connectorId;
private Map<String, Object> updateContent;
private MLUpdateConnectorRequest mlUpdateConnectorRequest;

@Mock
XContentParser parser;

@Before
public void setUp() {
MockitoAnnotations.openMocks(this);
this.connectorId = "test-connector_id";
this.updateContent = Map.of("description", "new description");
mlUpdateConnectorRequest = MLUpdateConnectorRequest.builder()
.connectorId(connectorId)
.updateContent(updateContent)
.build();
}

@Test
public void writeTo_Success() throws IOException {
BytesStreamOutput bytesStreamOutput = new BytesStreamOutput();
mlUpdateConnectorRequest.writeTo(bytesStreamOutput);
MLUpdateConnectorRequest parsedUpdateRequest = new MLUpdateConnectorRequest(bytesStreamOutput.bytes().streamInput());
assertEquals(connectorId, parsedUpdateRequest.getConnectorId());
assertEquals(updateContent, parsedUpdateRequest.getUpdateContent());
}

@Test
public void validate_Success() {
assertNull(mlUpdateConnectorRequest.validate());
}

@Test
public void validate_Exception_NullConnectorId() {
MLUpdateConnectorRequest updateConnectorRequest = MLUpdateConnectorRequest.builder().build();
Exception exception = updateConnectorRequest.validate();

assertEquals("Validation Failed: 1: ML connector id can't be null;", exception.getMessage());
}

@Test
public void parse_success() throws IOException {
RestRequest.Method method = RestRequest.Method.POST;
final Map<String, Object> updatefields = Map.of("version", "new version", "description", "new description");
when(parser.map()).thenReturn(updatefields);

MLUpdateConnectorRequest updateConnectorRequest = MLUpdateConnectorRequest.parse(parser, connectorId);
assertEquals(updateConnectorRequest.getConnectorId(), connectorId);
assertEquals(updateConnectorRequest.getUpdateContent(), updatefields);
}

@Test
public void fromActionRequest_Success() {
MLUpdateConnectorRequest mlUpdateConnectorRequest = MLUpdateConnectorRequest.builder()
.connectorId(connectorId)
.updateContent(updateContent)
.build();
assertSame(MLUpdateConnectorRequest.fromActionRequest(mlUpdateConnectorRequest), mlUpdateConnectorRequest);
}

@Test
public void fromActionRequest_Success_fromActionRequest() {
MLUpdateConnectorRequest mlUpdateConnectorRequest = MLUpdateConnectorRequest.builder()
.connectorId(connectorId)
.updateContent(updateContent)
.build();
ActionRequest actionRequest = new ActionRequest() {
@Override
public ActionRequestValidationException validate() {
return null;
}

@Override
public void writeTo(StreamOutput out) throws IOException {
mlUpdateConnectorRequest.writeTo(out);
}
};
MLUpdateConnectorRequest request = MLUpdateConnectorRequest.fromActionRequest(actionRequest);
assertNotSame(request, mlUpdateConnectorRequest);
assertEquals(mlUpdateConnectorRequest.getConnectorId(), request.getConnectorId());
assertEquals(mlUpdateConnectorRequest.getUpdateContent(), request.getUpdateContent());
}

@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();
}
};
MLUpdateConnectorRequest.fromActionRequest(actionRequest);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@
import org.opensearch.common.inject.Inject;
import org.opensearch.common.util.concurrent.ThreadContext;
import org.opensearch.core.action.ActionListener;
import org.opensearch.core.xcontent.NamedXContentRegistry;
import org.opensearch.ml.common.transport.connector.MLUpdateConnectorAction;
import org.opensearch.ml.common.transport.connector.MLUpdateConnectorRequest;
import org.opensearch.ml.helper.ConnectorAccessControlHelper;
Expand All @@ -32,7 +31,6 @@
@FieldDefaults(makeFinal = true, level = AccessLevel.PRIVATE)
public class UpdateConnectorTransportAction extends HandledTransportAction<ActionRequest, UpdateResponse> {
Client client;
NamedXContentRegistry xContentRegistry;

ConnectorAccessControlHelper connectorAccessControlHelper;

Expand All @@ -41,12 +39,10 @@ public UpdateConnectorTransportAction(
TransportService transportService,
ActionFilters actionFilters,
Client client,
NamedXContentRegistry xContentRegistry,
ConnectorAccessControlHelper connectorAccessControlHelper
) {
super(MLUpdateConnectorAction.NAME, transportService, actionFilters, MLUpdateConnectorRequest::new);
this.client = client;
this.xContentRegistry = xContentRegistry;
this.connectorAccessControlHelper = connectorAccessControlHelper;
}

Expand All @@ -69,11 +65,11 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Update
);
}
}, exception -> {
log.error("You don't have permission to update the connector for connector id: " + connectorId, exception);
log.error("Permission denied: Unable to update the connector with ID {}. Details: {}", connectorId, exception);
listener.onFailure(exception);
}));
} catch (Exception e) {
log.error("Failed to update ML connector " + connectorId, e);
log.error("Failed to update ML connector for connector id {}. Details {}:", connectorId, e);
listener.onFailure(e);
}
}
Expand All @@ -85,14 +81,14 @@ private ActionListener<UpdateResponse> getUpdateResponseListener(
) {
return ActionListener.runBefore(ActionListener.wrap(updateResponse -> {
if (updateResponse != null && updateResponse.getResult() != DocWriteResponse.Result.UPDATED) {
log.info("Connector id:{} failed update", connectorId);
log.info("Failed to update the connector with ID: {}", connectorId);
actionListener.onResponse(updateResponse);
return;
}
log.info("Completed Update Connector Request, connector id:{} updated", connectorId);
log.info("Successfully updated the connector with ID: {}", connectorId);
actionListener.onResponse(updateResponse);
}, exception -> {
log.error("Failed to update ML connector: " + connectorId, exception);
log.error("Failed to update ML connector with ID {}. Details: {}", connectorId, exception);
actionListener.onFailure(exception);
}), context::restore);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@
import java.util.List;
import java.util.Locale;

import org.apache.logging.log4j.util.Strings;
import org.opensearch.client.node.NodeClient;
import org.opensearch.core.xcontent.XContentParser;
import org.opensearch.ml.common.transport.connector.MLUpdateConnectorAction;
Expand Down Expand Up @@ -66,13 +65,10 @@ private MLUpdateConnectorRequest getRequest(RestRequest request) throws IOExcept
}

if (!request.hasContent()) {
throw new IOException("Update Connector request has empty body");
throw new IOException("Failed to update connector: Request body is empty");
}

String connectorId = getParameterId(request, PARAMETER_CONNECTOR_ID);
if (Strings.isBlank(connectorId)) {
throw new IOException("Update Connector request has no connector Id");
}

XContentParser parser = request.contentParser();
ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser);
Expand Down
Loading

0 comments on commit 9c70503

Please sign in to comment.