Skip to content

Commit f19614e

Browse files
authored
fixing security integ test (#3646)
Signed-off-by: Dhrubo Saha <dhrubo@amazon.com>
1 parent 17a9037 commit f19614e

File tree

1 file changed

+24
-3
lines changed

1 file changed

+24
-3
lines changed

plugin/src/test/java/org/opensearch/ml/rest/SecureMLRestIT.java

Lines changed: 24 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -409,10 +409,31 @@ public void testTrainWithReadOnlyMLAccess() throws IOException {
409409
}
410410

411411
public void testPredictWithReadOnlyMLAccess() throws IOException {
412-
exceptionRule.expect(ResponseException.class);
413-
exceptionRule.expectMessage("no permissions for [cluster:admin/opensearch/ml/predict]");
414412
KMeansParams kMeansParams = KMeansParams.builder().build();
415-
predict(mlReadOnlyClient, FunctionName.KMEANS, "modelId", irisIndex, kMeansParams, searchSourceBuilder, null);
413+
train(mlFullAccessClient, FunctionName.KMEANS, irisIndex, kMeansParams, searchSourceBuilder, trainResult -> {
414+
String modelId = (String) trainResult.get("model_id");
415+
assertNotNull(modelId);
416+
String status = (String) trainResult.get("status");
417+
assertEquals(MLTaskState.COMPLETED.name(), status);
418+
try {
419+
// Verify the model exists
420+
getModel(mlFullAccessClient, modelId, model -> {
421+
String algorithm = (String) model.get("algorithm");
422+
assertEquals(FunctionName.KMEANS.name(), algorithm);
423+
});
424+
425+
// Attempt prediction with read-only client and expect a permission error
426+
ResponseException exception = assertThrows(ResponseException.class, () -> {
427+
predict(mlReadOnlyClient, FunctionName.KMEANS, modelId, irisIndex, kMeansParams, searchSourceBuilder, null);
428+
});
429+
430+
// Verify the exception message
431+
assertTrue(exception.getMessage().contains("no permissions for [cluster:admin/opensearch/ml/predict]"));
432+
433+
} catch (IOException e) {
434+
fail("Unexpected IOException: " + e.getMessage());
435+
}
436+
}, false);
416437
}
417438

418439
public void testTrainAndPredictWithFullAccess() throws IOException {

0 commit comments

Comments
 (0)