From 9a9801faea29753b9016cace67da1cf00cd0893d Mon Sep 17 00:00:00 2001
From: HenryL27 <hmlindeman@yahoo.com>
Date: Wed, 29 Nov 2023 12:09:13 -0800
Subject: [PATCH] accept multiple failure if one of them is correct in kmeans
 prediction test

Signed-off-by: HenryL27 <hmlindeman@yahoo.com>
---
 .../action/prediction/PredictionITTests.java  | 22 +++++++++++++++----
 1 file changed, 18 insertions(+), 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
index 937167f00c..acadb4e7ce 100644
--- 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() {