Skip to content

Commit 9b40d25

Browse files
fixing security integ test (#3646) (#3780)
Signed-off-by: Dhrubo Saha <dhrubo@amazon.com> (cherry picked from commit f19614e) Co-authored-by: Dhrubo Saha <dhrubo@amazon.com>
1 parent c2e666c commit 9b40d25

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
@@ -408,10 +408,31 @@ public void testTrainWithReadOnlyMLAccess() throws IOException {
408408
}
409409

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

417438
public void testTrainAndPredictWithFullAccess() throws IOException {

0 commit comments

Comments
 (0)