From e97cfad38546651efb42b69070d4c621bd1c7e38 Mon Sep 17 00:00:00 2001
From: Robert Oskamp <robert.oskamp@elastic.co>
Date: Mon, 21 Jun 2021 16:11:28 +0200
Subject: [PATCH] [ML] Functional tests - stabilize module tests (#102708)

This PR stabilizes the `setupModule` API tests by increasing the timeout for the jobs to finish.
---
 .../apis/ml/modules/setup_module.ts           |  4 +--
 x-pack/test/functional/services/ml/api.ts     | 32 +++++++++++--------
 2 files changed, 20 insertions(+), 16 deletions(-)

diff --git a/x-pack/test/api_integration/apis/ml/modules/setup_module.ts b/x-pack/test/api_integration/apis/ml/modules/setup_module.ts
index 30dc31ef460c63..6011c38255cdc9 100644
--- a/x-pack/test/api_integration/apis/ml/modules/setup_module.ts
+++ b/x-pack/test/api_integration/apis/ml/modules/setup_module.ts
@@ -1146,8 +1146,8 @@ export default ({ getService }: FtrProviderContext) => {
             if (testData.requestBody.startDatafeed === true) {
               await ml.api.waitForADJobRecordCountToBePositive(job.jobId);
             }
-            await ml.api.waitForJobState(job.jobId, job.jobState);
-            await ml.api.waitForDatafeedState(datafeedId, job.datafeedState);
+            await ml.api.waitForDatafeedState(datafeedId, job.datafeedState, 4 * 60 * 1000);
+            await ml.api.waitForJobState(job.jobId, job.jobState, 4 * 60 * 1000);
 
             // model memory limit should be <= 99mb
             const {
diff --git a/x-pack/test/functional/services/ml/api.ts b/x-pack/test/functional/services/ml/api.ts
index 317f2dfe605143..728e3ff8fc8e6d 100644
--- a/x-pack/test/functional/services/ml/api.ts
+++ b/x-pack/test/functional/services/ml/api.ts
@@ -183,19 +183,19 @@ export function MachineLearningAPIProvider({ getService }: FtrProviderContext) {
       return jobStats;
     },
 
-    async waitForJobState(jobId: string, expectedJobState: JOB_STATE) {
-      await retry.waitForWithTimeout(
-        `job state to be ${expectedJobState}`,
-        2 * 60 * 1000,
-        async () => {
-          const state = await this.getJobState(jobId);
-          if (state === expectedJobState) {
-            return true;
-          } else {
-            throw new Error(`expected job state to be ${expectedJobState} but got ${state}`);
-          }
+    async waitForJobState(
+      jobId: string,
+      expectedJobState: JOB_STATE,
+      timeout: number = 2 * 60 * 1000
+    ) {
+      await retry.waitForWithTimeout(`job state to be ${expectedJobState}`, timeout, async () => {
+        const state = await this.getJobState(jobId);
+        if (state === expectedJobState) {
+          return true;
+        } else {
+          throw new Error(`expected job state to be ${expectedJobState} but got ${state}`);
         }
-      );
+      });
     },
 
     async getDatafeedState(datafeedId: string): Promise<DATAFEED_STATE> {
@@ -214,10 +214,14 @@ export function MachineLearningAPIProvider({ getService }: FtrProviderContext) {
       return state;
     },
 
-    async waitForDatafeedState(datafeedId: string, expectedDatafeedState: DATAFEED_STATE) {
+    async waitForDatafeedState(
+      datafeedId: string,
+      expectedDatafeedState: DATAFEED_STATE,
+      timeout: number = 2 * 60 * 1000
+    ) {
       await retry.waitForWithTimeout(
         `datafeed state to be ${expectedDatafeedState}`,
-        2 * 60 * 1000,
+        timeout,
         async () => {
           const state = await this.getDatafeedState(datafeedId);
           if (state === expectedDatafeedState) {