accept multiple failure if one of them is correct in kmeans predictio…

…n test Signed-off-by: HenryL27 <[email protected]>
HenryL27 · Dec 7, 2023 · 9a9801f · 9a9801f
1 parent a7df1df
commit 9a9801f
Showing 1 changed file with 18 additions and 4 deletions.
diff --git a/plugin/src/test/java/org/opensearch/ml/action/prediction/PredictionITTests.java b/plugin/src/test/java/org/opensearch/ml/action/prediction/PredictionITTests.java
@@ -14,6 +14,7 @@
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.rules.ExpectedException;
+import org.junit.runners.model.MultipleFailureException;
 import org.opensearch.action.ActionRequestValidationException;
 import org.opensearch.common.action.ActionFuture;
 import org.opensearch.ml.action.MLCommonsIntegTestCase;
@@ -89,14 +90,27 @@ public void testPredictionWithoutDataset_KMeans() {
         predictionFuture.actionGet();
     }
 
-    public void testPredictionWithEmptyDataset_KMeans() {
-        exceptionRule.expect(IllegalArgumentException.class);
-        exceptionRule.expectMessage("No document found");
+    public void testPredictionWithEmptyDataset_KMeans() throws Exception {
         MLInputDataset emptySearchInputDataset = emptyQueryInputDataSet(irisIndexName);
         MLInput mlInput = MLInput.builder().algorithm(FunctionName.KMEANS).inputDataset(emptySearchInputDataset).build();
         MLPredictionTaskRequest predictionRequest = new MLPredictionTaskRequest(kMeansModelId, mlInput, null);
         ActionFuture<MLTaskResponse> predictionFuture = client().execute(MLPredictionTaskAction.INSTANCE, predictionRequest);
-        predictionFuture.actionGet();
+        Exception e = assertThrows(Exception.class, () -> predictionFuture.actionGet());
+        boolean foundIllegalArgument = false;
+        if (e instanceof MultipleFailureException) {
+            for (Throwable t : ((MultipleFailureException) e).getFailures()) {
+                if (t instanceof IllegalArgumentException) {
+                    e = (Exception) t;
+                    foundIllegalArgument = true;
+                    break;
+                }
+            }
+            if (!foundIllegalArgument) {
+                throw e;
+            }
+        }
+        assert (e instanceof IllegalArgumentException);
+        assert (e.getMessage().equals("No document found"));
     }
 
     public void testPredictionWithSearchInput_LogisticRegression() {