[ML] memory limit test fix, waiting until task is failed before stopping (#61576)

x-pack/plugin/ml/qa/native-multi-node-tests/src/test/java/org/elasticsearch/xpack/ml/integration/ClassificationIT.java

@@ -790,6 +790,7 @@ public class ClassificationIT extends MlNativeDataFrameAnalyticsIntegTestCase {
         putAnalytics(config);
         // Shouldn't throw
         startAnalytics(jobId);
+        waitUntilAnalyticsIsFailed(jobId);
         // It could be marked as failed...
         forceStopAnalytics(jobId);
         waitUntilAnalyticsIsStopped(jobId);

x-pack/plugin/ml/qa/native-multi-node-tests/src/test/java/org/elasticsearch/xpack/ml/integration/MlNativeDataFrameAnalyticsIntegTestCase.java

@@ -157,6 +157,10 @@ abstract class MlNativeDataFrameAnalyticsIntegTestCase extends MlNativeIntegTest
         assertBusy(() -> assertIsStopped(id), waitTime.getMillis(), TimeUnit.MILLISECONDS);
     }
 
+    protected void waitUntilAnalyticsIsFailed(String id) throws Exception {
+        assertBusy(() -> assertIsFailed(id), TimeValue.timeValueSeconds(30).millis(), TimeUnit.MILLISECONDS);
+    }
+
     protected List<DataFrameAnalyticsConfig> getAnalytics(String id) {
         GetDataFrameAnalyticsAction.Request request = new GetDataFrameAnalyticsAction.Request(id);
         return client().execute(GetDataFrameAnalyticsAction.INSTANCE, request).actionGet().getResources().results();
@@ -207,6 +211,11 @@ abstract class MlNativeDataFrameAnalyticsIntegTestCase extends MlNativeIntegTest
         assertThat("Stats were: " + Strings.toString(stats), stats.getState(), equalTo(DataFrameAnalyticsState.STOPPED));
     }
 
+    protected void assertIsFailed(String id) {
+        GetDataFrameAnalyticsStatsAction.Response.Stats stats = getAnalyticsStats(id);
+        assertThat("Stats were: " + Strings.toString(stats), stats.getState(), equalTo(DataFrameAnalyticsState.FAILED));
+    }
+
     protected void assertProgressIsZero(String id) {
         List<PhaseProgress> progress = getProgress(id);
         assertThat("progress is not all zero: " + progress,

x-pack/plugin/ml/qa/native-multi-node-tests/src/test/java/org/elasticsearch/xpack/ml/integration/RunDataFrameAnalyticsIT.java

@@ -519,7 +519,7 @@ public class RunDataFrameAnalyticsIT extends MlNativeDataFrameAnalyticsIntegTest
         assertIsStopped(id);
         //should not throw
         startAnalytics(id);
-        // Might have been marked as failed
+        waitUntilAnalyticsIsFailed(id);
         forceStopAnalytics(id);
         waitUntilAnalyticsIsStopped(id);
     }