[ML] Fix flaky ForecastIT#testOverflowToDisk (#137452)

Author: valeriy42

muted-tests.yml

@@ -94,9 +94,6 @@ tests:
 - class: org.elasticsearch.smoketest.MlWithSecurityIT
   method: test {yaml=ml/sparse_vector_search/Test sparse_vector search with query vector and pruning config}
   issue: https://github.com/elastic/elasticsearch/issues/119548
-- class: org.elasticsearch.xpack.ml.integration.ForecastIT
-  method: testOverflowToDisk
-  issue: https://github.com/elastic/elasticsearch/issues/117740
 - class: org.elasticsearch.multi_cluster.MultiClusterYamlTestSuiteIT
   issue: https://github.com/elastic/elasticsearch/issues/119983
 - class: org.elasticsearch.xpack.test.rest.XPackRestIT

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

@@ -246,7 +246,19 @@ protected String forecast(String jobId, TimeValue duration, TimeValue expiresIn,
 
     protected void waitForecastToFinish(String jobId, String forecastId) throws Exception {
         // Forecasts can take an eternity to complete in the FIPS JVM
-        waitForecastStatus(inFipsJvm() ? 300 : 90, jobId, forecastId, ForecastRequestStats.ForecastRequestStatus.FINISHED);
+        int timeoutSeconds = inFipsJvm() ? 300 : 90;
+        // First wait for the forecast document to exist and be in a non-terminal state
+        // This handles the race condition where the document may be SCHEDULED or STARTED initially
+        waitForecastStatus(
+            timeoutSeconds,
+            jobId,
+            forecastId,
+            ForecastRequestStats.ForecastRequestStatus.SCHEDULED,
+            ForecastRequestStats.ForecastRequestStatus.STARTED,
+            ForecastRequestStats.ForecastRequestStatus.FINISHED
+        );
+        // Then wait specifically for FINISHED status
+        waitForecastStatus(timeoutSeconds, jobId, forecastId, ForecastRequestStats.ForecastRequestStatus.FINISHED);
     }
 
     protected void waitForecastStatus(String jobId, String forecastId, ForecastRequestStats.ForecastRequestStatus... status)
@@ -261,6 +273,10 @@ protected void waitForecastStatus(
         ForecastRequestStats.ForecastRequestStatus... status
     ) throws Exception {
         assertBusy(() -> {
+            // Refresh the index to ensure recently indexed forecast stats documents are visible
+            indicesAdmin().prepareRefresh(AnomalyDetectorsIndex.jobResultsAliasedName(jobId))
+                .setIndicesOptions(IndicesOptions.LENIENT_EXPAND_OPEN_HIDDEN)
+                .get();
             ForecastRequestStats forecastRequestStats = getForecastStats(jobId, forecastId);
             assertThat(forecastRequestStats, is(notNullValue()));
             assertThat(forecastRequestStats.getStatus(), in(status));