Skip to content

Commit

Permalink
remove return_content field and add more tests
Browse files Browse the repository at this point in the history 
Signed-off-by: Bhavana Ramaram <[email protected]>
  • Loading branch information
@rbhavna
rbhavna committed Nov 24, 2023
1 parent 5ac3c9f commit d0a7491
Show file tree
Hide file tree
Showing 6 changed files with 163 additions and 24 deletions.
6 changes: 1 addition & 5 deletions .../src/main/java/org/opensearch/ml/common/transport/model_group/MLModelGroupGetRequest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,25 +30,21 @@
public class MLModelGroupGetRequest extends ActionRequest {

String modelGroupId;
boolean returnContent;

@Builder
public MLModelGroupGetRequest(String modelGroupId, boolean returnContent) {
public MLModelGroupGetRequest(String modelGroupId) {
this.modelGroupId = modelGroupId;
this.returnContent = returnContent;
}

public MLModelGroupGetRequest(StreamInput in) throws IOException {
super(in);
this.modelGroupId = in.readString();
this.returnContent = in.readBoolean();
}

@Override
public void writeTo(StreamOutput out) throws IOException {
super.writeTo(out);
out.writeString(this.modelGroupId);
out.writeBoolean(returnContent);
}

@Override
Expand Down
11 changes: 4 additions & 7 deletions plugin/src/main/java/org/opensearch/ml/action/model_group/GetModelGroupTransportAction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@
import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
import static org.opensearch.ml.common.CommonValue.ML_MODEL_GROUP_INDEX;
import static org.opensearch.ml.utils.MLNodeUtils.createXContentParserFromRegistry;
import static org.opensearch.ml.utils.RestActionUtils.getFetchSourceContext;

import org.opensearch.OpenSearchStatusException;
import org.opensearch.action.ActionRequest;
Expand All @@ -27,13 +26,11 @@
import org.opensearch.index.IndexNotFoundException;
import org.opensearch.ml.common.MLModelGroup;
import org.opensearch.ml.common.exception.MLResourceNotFoundException;
import org.opensearch.ml.common.exception.MLValidationException;
import org.opensearch.ml.common.transport.model_group.MLModelGroupGetAction;
import org.opensearch.ml.common.transport.model_group.MLModelGroupGetRequest;
import org.opensearch.ml.common.transport.model_group.MLModelGroupGetResponse;
import org.opensearch.ml.helper.ModelAccessControlHelper;
import org.opensearch.ml.utils.RestActionUtils;
import org.opensearch.search.fetch.subphase.FetchSourceContext;
import org.opensearch.tasks.Task;
import org.opensearch.transport.TransportService;

Expand Down Expand Up @@ -71,8 +68,7 @@ public GetModelGroupTransportAction(
protected void doExecute(Task task, ActionRequest request, ActionListener<MLModelGroupGetResponse> actionListener) {
MLModelGroupGetRequest mlModelGroupGetRequest = MLModelGroupGetRequest.fromActionRequest(request);
String modelGroupId = mlModelGroupGetRequest.getModelGroupId();
FetchSourceContext fetchSourceContext = getFetchSourceContext(mlModelGroupGetRequest.isReturnContent());
GetRequest getRequest = new GetRequest(ML_MODEL_GROUP_INDEX).id(modelGroupId).fetchSourceContext(fetchSourceContext);
GetRequest getRequest = new GetRequest(ML_MODEL_GROUP_INDEX).id(modelGroupId);
User user = RestActionUtils.getUserContext(client);

try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
Expand All @@ -87,8 +83,9 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLMode
if (!access) {
wrappedListener
.onFailure(
new MLValidationException(
"User doesn't have privilege to perform this operation on this model group"
new OpenSearchStatusException(
"User doesn't have privilege to perform this operation on this model group",
RestStatus.FORBIDDEN
)
);
} else {
Expand Down
4 changes: 1 addition & 3 deletions plugin/src/main/java/org/opensearch/ml/rest/RestMLGetModelGroupAction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@
import static org.opensearch.ml.plugin.MachineLearningPlugin.ML_BASE_URI;
import static org.opensearch.ml.utils.RestActionUtils.PARAMETER_MODEL_GROUP_ID;
import static org.opensearch.ml.utils.RestActionUtils.getParameterId;
import static org.opensearch.ml.utils.RestActionUtils.returnContent;

import java.io.IOException;
import java.util.List;
Expand Down Expand Up @@ -60,8 +59,7 @@ public RestChannelConsumer prepareRequest(RestRequest request, NodeClient client
@VisibleForTesting
MLModelGroupGetRequest getRequest(RestRequest request) throws IOException {
String modelGroupId = getParameterId(request, PARAMETER_MODEL_GROUP_ID);
boolean returnContent = returnContent(request);

return new MLModelGroupGetRequest(modelGroupId, returnContent);
return new MLModelGroupGetRequest(modelGroupId);
}
}
2 changes: 1 addition & 1 deletion plugin/src/test/java/org/opensearch/ml/action/MLCommonsIntegTestCase.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -391,7 +391,7 @@ public MLModel getModel(String modelId) {
}

public MLModelGroup getModelGroup(String modelGroupId) {
MLModelGroupGetRequest getRequest = new MLModelGroupGetRequest(modelGroupId, false);
MLModelGroupGetRequest getRequest = new MLModelGroupGetRequest(modelGroupId);
MLModelGroupGetResponse response = client().execute(MLModelGroupGetAction.INSTANCE, getRequest).actionGet(5000);
return response.getMlModelGroup();
}
Expand Down
11 changes: 10 additions & 1 deletion plugin/src/test/java/org/opensearch/ml/rest/MLCommonsRestTestCase.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,11 @@
import org.apache.hc.client5.http.impl.nio.PoolingAsyncClientConnectionManagerBuilder;
import org.apache.hc.client5.http.ssl.ClientTlsStrategyBuilder;
import org.apache.hc.client5.http.ssl.NoopHostnameVerifier;
import org.apache.hc.core5.http.*;
import org.apache.hc.core5.http.Header;
import org.apache.hc.core5.http.HttpEntity;
import org.apache.hc.core5.http.HttpHeaders;
import org.apache.hc.core5.http.HttpHost;
import org.apache.hc.core5.http.ParseException;
import org.apache.hc.core5.http.io.entity.EntityUtils;
import org.apache.hc.core5.http.message.BasicHeader;
import org.apache.hc.core5.http.nio.ssl.TlsStrategy;
Expand Down Expand Up @@ -718,6 +722,11 @@ public void deleteModelGroup(RestClient client, String modelGroupId, Consumer<Ma
verifyResponse(function, response);
}

public void getModelGroup(RestClient client, String modelIGroupd, Consumer<Map<String, Object>> function) throws IOException {
Response response = TestHelper.makeRequest(client, "GET", "/_plugins/_ml/model_groups/" + modelIGroupd, null, "", null);
verifyResponse(function, response);
}

public void searchModelGroups(RestClient client, String query, Consumer<Map<String, Object>> function) throws IOException {
Response response = TestHelper.makeRequest(client, "GET", "/_plugins/_ml/model_groups/_search", null, query, null);
verifyResponse(function, response);
Expand Down
153 changes: 146 additions & 7 deletions plugin/src/test/java/org/opensearch/ml/rest/MLModelGroupRestIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,11 +7,8 @@

package org.opensearch.ml.rest;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Map;

import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableList;
import org.apache.hc.core5.http.HttpHeaders;
import org.apache.hc.core5.http.HttpHost;
import org.apache.hc.core5.http.message.BasicHeader;
Expand All @@ -28,8 +25,10 @@
import org.opensearch.ml.common.transport.model_group.MLUpdateModelGroupInput;
import org.opensearch.ml.utils.TestHelper;

import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableList;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Map;

public class MLModelGroupRestIT extends MLCommonsRestTestCase {

Expand Down Expand Up @@ -1212,4 +1211,144 @@ public void test_search_MatchAllQuery_For_ModelGroups() throws IOException {
}
}

public void test_get_modelGroup() throws IOException {
mlRegisterModelGroupInput = createRegisterModelGroupInput("testModelGroup1", Arrays.asList("IT"), AccessMode.RESTRICTED, null);
registerModelGroup(user1Client, TestHelper.toJsonString(mlRegisterModelGroupInput), registerModelGroupResult -> {
String modelGroupId1 = (String) registerModelGroupResult.get("model_group_id");
assertTrue(registerModelGroupResult.containsKey("model_group_id"));
try {
// User2 successfully gets model group since user2 has IT backend role
getModelGroup(
user2Client,
modelGroupId1,
getModelGroupResult -> { assertTrue(getModelGroupResult.containsKey("model_group_id")); }
);

// Admin successfully gets model group
getModelGroup(
client(),
modelGroupId1,
getModelGroupResult -> { assertTrue(getModelGroupResult.containsKey("model_group_id")); }
);
} catch (IOException e) {
assertNull(e);
}
// User2 fails to get model group
try {
getModelGroup(user3Client, modelGroupId, null);
} catch (Exception e) {
assertEquals(ResponseException.class, e.getClass());
assertTrue(
Throwables
.getStackTraceAsString(e)
.contains("User doesn't have privilege to perform this operation on this model group")
);
}
});

mlRegisterModelGroupInput = createRegisterModelGroupInput("testModelGroup2", null, AccessMode.PUBLIC, null);
registerModelGroup(user2Client, TestHelper.toJsonString(mlRegisterModelGroupInput), registerModelGroupResult -> {
String modelGroupId2 = (String) registerModelGroupResult.get("model_group_id");
assertTrue(registerModelGroupResult.containsKey("model_group_id"));
try {
// User1 successfully gets model group since user2 has IT backend role
getModelGroup(
user1Client,
modelGroupId2,
getModelGroupResult -> { assertTrue(getModelGroupResult.containsKey("model_group_id")); }
);

// User3 successfully gets model group
getModelGroup(
user3Client,
modelGroupId2,
getModelGroupResult -> { assertTrue(getModelGroupResult.containsKey("model_group_id")); }
);

// User4 successfully gets model group
getModelGroup(
user4Client,
modelGroupId2,
getModelGroupResult -> { assertTrue(getModelGroupResult.containsKey("model_group_id")); }
);
} catch (IOException e) {
assertNull(e);
}
});

mlRegisterModelGroupInput = createRegisterModelGroupInput("testModelGroup3", null, AccessMode.PRIVATE, null);
registerModelGroup(user3Client, TestHelper.toJsonString(mlRegisterModelGroupInput), registerModelGroupResult -> {
String modelGroupId3 = (String) registerModelGroupResult.get("model_group_id");
assertTrue(registerModelGroupResult.containsKey("model_group_id"));
try {
// User3 successfully gets model group since user2 has IT backend role
getModelGroup(
user3Client,
modelGroupId3,
getModelGroupResult -> { assertTrue(getModelGroupResult.containsKey("model_group_id")); }
);

// Admin successfully gets model group
getModelGroup(
client(),
modelGroupId3,
getModelGroupResult -> { assertTrue(getModelGroupResult.containsKey("model_group_id")); }
);
} catch (IOException e) {
assertNull(e);
}
// User2 fails to get model group
try {
getModelGroup(user2Client, modelGroupId3, null);
} catch (Exception e) {
assertEquals(ResponseException.class, e.getClass());
assertTrue(
Throwables
.getStackTraceAsString(e)
.contains("User doesn't have privilege to perform this operation on this model group")
);
}
});

mlRegisterModelGroupInput = createRegisterModelGroupInput("testModelGroup4", null, null, null);
registerModelGroup(client(), TestHelper.toJsonString(mlRegisterModelGroupInput), registerModelGroupResult -> {
String modelGroupId4 = (String) registerModelGroupResult.get("model_group_id");
assertTrue(registerModelGroupResult.containsKey("model_group_id"));
try {
// Admin successfully gets model group
getModelGroup(
client(),
modelGroupId4,
getModelGroupResult -> { assertTrue(getModelGroupResult.containsKey("model_group_id")); }
);
} catch (IOException e) {
assertNull(e);
}

// User1 fails to get model group
try {
getModelGroup(user1Client, modelGroupId4, null);
} catch (Exception e) {
assertEquals(ResponseException.class, e.getClass());
assertTrue(
Throwables
.getStackTraceAsString(e)
.contains("User doesn't have privilege to perform this operation on this model group")
);
}

// User2 fails to get model group
try {
getModelGroup(user2Client, modelGroupId4, null);
} catch (Exception e) {
assertEquals(ResponseException.class, e.getClass());
assertTrue(
Throwables
.getStackTraceAsString(e)
.contains("User doesn't have privilege to perform this operation on this model group")
);
}
});
}

}

0 comments on commit d0a7491

Please sign in to comment.