From aa53e7177a2fc6de125a64077f8f6ea6a7e5d750 Mon Sep 17 00:00:00 2001
From: Martijn van Groningen <martijn.v.groningen@gmail.com>
Date: Thu, 24 Nov 2016 11:49:44 +0100
Subject: [PATCH] wait for the allocation to be added before sending data

Original commit: elastic/x-pack-elasticsearch@504a75a2f291be3b9f041020e423ec9bb783a4ce
---
 .../xpack/prelert/action/GetJobsAction.java   | 22 ++++++++++++++++++-
 .../prelert/integration/TooManyJobsIT.java    | 20 ++++++++++++++++-
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/elasticsearch/src/main/java/org/elasticsearch/xpack/prelert/action/GetJobsAction.java b/elasticsearch/src/main/java/org/elasticsearch/xpack/prelert/action/GetJobsAction.java
index 1ee22cbb71a..fa430738844 100644
--- a/elasticsearch/src/main/java/org/elasticsearch/xpack/prelert/action/GetJobsAction.java
+++ b/elasticsearch/src/main/java/org/elasticsearch/xpack/prelert/action/GetJobsAction.java
@@ -245,7 +245,7 @@ public class GetJobsAction extends Action<GetJobsAction.Request, GetJobsAction.R
 
     public static class Response extends ActionResponse implements StatusToXContent {
 
-        static class JobInfo implements ToXContent, Writeable {
+        public static class JobInfo implements ToXContent, Writeable {
             @Nullable
             private Job jobConfig;
             @Nullable
@@ -276,6 +276,26 @@ public class GetJobsAction extends Action<GetJobsAction.Request, GetJobsAction.R
                 status = in.readOptionalWriteable(JobStatus::fromStream);
             }
 
+            public Job getJobConfig() {
+                return jobConfig;
+            }
+
+            public DataCounts getDataCounts() {
+                return dataCounts;
+            }
+
+            public ModelSizeStats getModelSizeStats() {
+                return modelSizeStats;
+            }
+
+            public SchedulerState getSchedulerState() {
+                return schedulerState;
+            }
+
+            public JobStatus getStatus() {
+                return status;
+            }
+
             @Override
             public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
                 builder.startObject();
diff --git a/elasticsearch/src/test/java/org/elasticsearch/xpack/prelert/integration/TooManyJobsIT.java b/elasticsearch/src/test/java/org/elasticsearch/xpack/prelert/integration/TooManyJobsIT.java
index 05ba59ba103..5dccb72bddc 100644
--- a/elasticsearch/src/test/java/org/elasticsearch/xpack/prelert/integration/TooManyJobsIT.java
+++ b/elasticsearch/src/test/java/org/elasticsearch/xpack/prelert/integration/TooManyJobsIT.java
@@ -13,6 +13,7 @@ import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.test.ESIntegTestCase;
 import org.elasticsearch.xpack.prelert.PrelertPlugin;
+import org.elasticsearch.xpack.prelert.action.GetJobsAction;
 import org.elasticsearch.xpack.prelert.action.PostDataAction;
 import org.elasticsearch.xpack.prelert.action.PutJobAction;
 import org.elasticsearch.xpack.prelert.action.ScheduledJobsIT;
@@ -20,10 +21,12 @@ import org.elasticsearch.xpack.prelert.job.AnalysisConfig;
 import org.elasticsearch.xpack.prelert.job.DataDescription;
 import org.elasticsearch.xpack.prelert.job.Detector;
 import org.elasticsearch.xpack.prelert.job.Job;
+import org.elasticsearch.xpack.prelert.job.JobStatus;
 import org.junit.After;
 
 import java.util.Collection;
 import java.util.Collections;
+import java.util.concurrent.ExecutionException;
 
 @ESIntegTestCase.ClusterScope(numDataNodes = 1)
 public class TooManyJobsIT extends ESIntegTestCase {
@@ -51,6 +54,19 @@ public class TooManyJobsIT extends ESIntegTestCase {
             PutJobAction.Request putJobRequest = new PutJobAction.Request(job.build(true));
             PutJobAction.Response putJobResponse = client().execute(PutJobAction.INSTANCE, putJobRequest).get();
             assertTrue(putJobResponse.isAcknowledged());
+            assertBusy(() -> {
+                try {
+                    GetJobsAction.Request getJobRequest = new GetJobsAction.Request();
+                    getJobRequest.setJobId(job.getId());
+                    getJobRequest.status(true);
+                    GetJobsAction.Response response = client().execute(GetJobsAction.INSTANCE, getJobRequest).get();
+                    GetJobsAction.Response.JobInfo jobInfo = response.getResponse().results().get(0);
+                    assertNotNull(jobInfo);
+                    assertEquals(JobStatus.CLOSED, jobInfo.getStatus());
+                } catch (Exception e) {
+                    fail("failure " + e.getMessage());
+                }
+            });
 
             // triggers creating autodetect process:
             PostDataAction.Request postDataRequest = new PostDataAction.Request(job.getId());
@@ -60,8 +76,10 @@ public class TooManyJobsIT extends ESIntegTestCase {
                 assertEquals(1, postDataResponse.getDataCounts().getInputRecordCount());
                 logger.info("Posted data {} times", i);
             } catch (Exception e) {
-                logger.info("Handling exception", e);
                 Throwable cause = ExceptionsHelper.unwrapCause(e.getCause());
+                if (ElasticsearchStatusException.class.equals(cause.getClass()) == false) {
+                    logger.warn("Unexpected cause", e);
+                }
                 assertEquals(ElasticsearchStatusException.class, cause.getClass());
                 assertEquals(RestStatus.TOO_MANY_REQUESTS, ((ElasticsearchStatusException) cause).status());
                 logger.info("good news everybody --> reached threadpool capacity after starting {}th analytical process", i);