[ML] Get ForecastRequestStats doc in RestoreModelSnapshotIT (#31973)

Author: davidkyle

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

@@ -6,6 +6,7 @@
 package org.elasticsearch.xpack.ml.integration;
 
 import org.elasticsearch.action.admin.cluster.node.tasks.list.ListTasksRequest;
+import org.elasticsearch.action.get.GetResponse;
 import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.client.Client;
 import org.elasticsearch.cluster.ClusterModule;
@@ -342,21 +343,17 @@ protected void waitForecastToFinish(String jobId, String forecastId) throws Exce
     }
 
     protected ForecastRequestStats getForecastStats(String jobId, String forecastId) {
-        SearchResponse searchResponse = client().prepareSearch(AnomalyDetectorsIndex.jobResultsAliasedName(jobId))
-                .setQuery(QueryBuilders.boolQuery()
-                        .filter(QueryBuilders.termQuery(Result.RESULT_TYPE.getPreferredName(), ForecastRequestStats.RESULT_TYPE_VALUE))
-                        .filter(QueryBuilders.termQuery(Job.ID.getPreferredName(), jobId))
-                        .filter(QueryBuilders.termQuery(ForecastRequestStats.FORECAST_ID.getPreferredName(), forecastId)))
+        GetResponse getResponse = client().prepareGet()
+                .setIndex(AnomalyDetectorsIndex.jobResultsAliasedName(jobId))
+                .setId(ForecastRequestStats.documentId(jobId, forecastId))
                 .execute().actionGet();
-        SearchHits hits = searchResponse.getHits();
-        if (hits.getTotalHits() == 0) {
+
+        if (getResponse.isExists() == false) {
             return null;
         }
-        assertThat(hits.getTotalHits(), equalTo(1L));
-        try {
-            XContentParser parser = XContentFactory.xContent(XContentType.JSON).createParser(
+        try (XContentParser parser = XContentFactory.xContent(XContentType.JSON).createParser(
                     NamedXContentRegistry.EMPTY, DeprecationHandler.THROW_UNSUPPORTED_OPERATION,
-                    hits.getHits()[0].getSourceRef().streamInput());
+                    getResponse.getSourceAsBytesRef().streamInput())) {
             return ForecastRequestStats.STRICT_PARSER.apply(parser, null);
         } catch (IOException e) {
             throw new IllegalStateException(e);

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

@@ -23,7 +23,11 @@
 import java.util.Map;
 import java.util.stream.Collectors;
 
+import static org.hamcrest.Matchers.anyOf;
+import static org.hamcrest.Matchers.empty;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.greaterThan;
+import static org.hamcrest.Matchers.nullValue;
 
 /**
  * This test aims to catch regressions where,
@@ -69,7 +73,9 @@ public void test() throws Exception {
         String forecastId = forecast(job.getId(), TimeValue.timeValueHours(3), null);
         waitForecastToFinish(job.getId(), forecastId);
         ForecastRequestStats forecastStats = getForecastStats(job.getId(), forecastId);
-        assertThat(forecastStats.getStatus(), equalTo(ForecastRequestStats.ForecastRequestStatus.FINISHED));
+        assertThat(forecastStats.getMessages(), anyOf(nullValue(), empty()));
+        assertThat(forecastStats.getMemoryUsage(), greaterThan(0L));
+        assertEquals(forecastStats.getRecordCount(), 3L);
 
         closeJob(job.getId());