Fix the tests

Author: davidkyle

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/job/retention/AbstractExpiredJobDataRemover.java

@@ -70,8 +70,15 @@ private void removeData(WrappedBatchedJobsIterator jobIterator, ActionListener<B
         }
 
         calcCutoffEpochMs(job.getId(), retentionDays, ActionListener.wrap(
-                cutoffEpochMs -> removeDataBefore(job, cutoffEpochMs,
-                        ActionListener.wrap(response -> removeData(jobIterator, listener, isTimedOutSupplier), listener::onFailure)),
+                cutoffEpochMs -> {
+                    if (cutoffEpochMs == null) {
+                        removeData(jobIterator, listener, isTimedOutSupplier);
+                    } else {
+                        removeDataBefore(job, cutoffEpochMs, ActionListener.wrap(
+                                response -> removeData(jobIterator, listener, isTimedOutSupplier),
+                                listener::onFailure));
+                    }
+                },
                 listener::onFailure
         ));
     }

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/job/retention/ExpiredModelSnapshotsRemover.java

@@ -30,8 +30,6 @@
 import org.elasticsearch.xpack.core.ml.job.persistence.ElasticsearchMappings;
 import org.elasticsearch.xpack.core.ml.job.process.autodetect.state.ModelSnapshot;
 import org.elasticsearch.xpack.core.ml.job.process.autodetect.state.ModelSnapshotField;
-import org.elasticsearch.xpack.core.ml.job.results.Bucket;
-import org.elasticsearch.xpack.core.ml.job.results.Result;
 import org.elasticsearch.xpack.ml.MachineLearning;
 import org.elasticsearch.xpack.ml.utils.MlIndicesUtils;
 import org.elasticsearch.xpack.ml.utils.VolatileCursorIterator;
@@ -84,20 +82,25 @@ void calcCutoffEpochMs(String jobId, long retentionDays, ActionListener<Long> li
 
         latestSnapshotTimeStamp(jobId, ActionListener.wrap(
                 latestTime -> {
-                    long cutoff = latestTime - new TimeValue(retentionDays, TimeUnit.DAYS).getMillis();
-                    threadedActionListener.onResponse(cutoff);
+                    if (latestTime == null) {
+                        threadedActionListener.onResponse(null);
+                    } else {
+                        long cutoff = latestTime - new TimeValue(retentionDays, TimeUnit.DAYS).getMillis();
+                        threadedActionListener.onResponse(cutoff);
+                    }
                 },
                 listener::onFailure
         ));
     }
 
     private void latestSnapshotTimeStamp(String jobId, ActionListener<Long> listener) {
         SortBuilder<?> sortBuilder = new FieldSortBuilder(ModelSnapshot.TIMESTAMP.getPreferredName()).order(SortOrder.DESC);
-        QueryBuilder bucketType = QueryBuilders.termQuery(Result.RESULT_TYPE.getPreferredName(), Bucket.RESULT_TYPE_VALUE);
+        QueryBuilder snapshotQuery = QueryBuilders.boolQuery()
+                .filter(QueryBuilders.existsQuery(ModelSnapshot.SNAPSHOT_DOC_COUNT.getPreferredName()));
 
         SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
         searchSourceBuilder.sort(sortBuilder);
-        searchSourceBuilder.query(bucketType);
+        searchSourceBuilder.query(snapshotQuery);
         searchSourceBuilder.size(1);
         searchSourceBuilder.trackTotalHits(false);
 

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/job/retention/ExpiredResultsRemover.java

@@ -136,8 +136,12 @@ void calcCutoffEpochMs(String jobId, long retentionDays, ActionListener<Long> li
                 MachineLearning.UTILITY_THREAD_POOL_NAME, listener, false);
         latestBucketTime(jobId, ActionListener.wrap(
                 latestTime -> {
-                    long cutoff = latestTime - new TimeValue(retentionDays, TimeUnit.DAYS).getMillis();
-                    threadedActionListener.onResponse(cutoff);
+                    if (latestTime == null) {
+                        threadedActionListener.onResponse(null);
+                    } else {
+                        long cutoff = latestTime - new TimeValue(retentionDays, TimeUnit.DAYS).getMillis();
+                        threadedActionListener.onResponse(cutoff);
+                    }
                 },
                 listener::onFailure
         ));

x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/job/retention/AbstractExpiredJobDataRemoverTests.java

@@ -82,10 +82,10 @@ static SearchResponse createSearchResponse(List<? extends ToXContent> toXContent
     static void givenJobs(Client client, List<Job> jobs) throws IOException {
         SearchResponse response = AbstractExpiredJobDataRemoverTests.createSearchResponse(jobs);
 
-        doAnswer(invocationOnMock -> {
-            ActionListener<SearchResponse> listener = (ActionListener<SearchResponse>) invocationOnMock.getArguments()[2];
-            listener.onResponse(response);
-            return null;
+            doAnswer(invocationOnMock -> {
+                ActionListener<SearchResponse> listener = (ActionListener<SearchResponse>) invocationOnMock.getArguments()[2];
+                listener.onResponse(response);
+                return null;
         }).when(client).execute(eq(SearchAction.INSTANCE), any(), any());
     }