[kbn-test] retry Cloud token fetching (#186626)

## Summary

The internal request to security service has no retry logic, so
sometimes Cloud API call to fetch token fails on timeout (180s) with 503
code. Still the success rate in QA is high, so decision is to to retry
fetching token on client side.

PR adds few changes:
- try to fetch valid Cloud token with 3 attempts with 15 sec delay
before next attempt
- set request timeout explicitly to 60s to not "hang" for 180s

---------

Co-authored-by: Tre' Seymour <wayne.seymour@elastic.co>

packages/kbn-test/src/auth/saml_auth.test.ts

@@ -42,8 +42,11 @@ const mockGetOnce = (mockedUrl: string, response: any) => {
 
 describe('saml_auth', () => {
   describe('createCloudSession', () => {
+    afterEach(() => {
+      axiosRequestMock.mockClear();
+    });
     test('returns token value', async () => {
-      mockRequestOnce('/api/v1/saas/auth/_login', { data: { token: 'mocked_token' } });
+      mockRequestOnce('/api/v1/saas/auth/_login', { data: { token: 'mocked_token' }, status: 200 });
 
       const sessionToken = await createCloudSession({
         hostname: 'cloud',
@@ -52,23 +55,124 @@ describe('saml_auth', () => {
         log,
       });
       expect(sessionToken).toBe('mocked_token');
+      expect(axiosRequestMock).toBeCalledTimes(1);
     });
 
-    test('throws error when response has no token', async () => {
-      mockRequestOnce('/api/v1/saas/auth/_login', { data: { message: 'no token' } });
+    test('retries until response has the token value', async () => {
+      let callCount = 0;
+      axiosRequestMock.mockImplementation((config: AxiosRequestConfig) => {
+        if (config.url?.endsWith('/api/v1/saas/auth/_login')) {
+          callCount += 1;
+          if (callCount !== 3) {
+            return Promise.resolve({ data: { message: 'no token' }, status: 503 });
+          } else {
+            return Promise.resolve({
+              data: { token: 'mocked_token' },
+              status: 200,
+            });
+          }
+        }
+        return Promise.reject(new Error(`Unexpected URL: ${config.url}`));
+      });
 
-      await expect(
-        createCloudSession({
+      const sessionToken = await createCloudSession(
+        {
           hostname: 'cloud',
           email: 'viewer@elastic.co',
           password: 'changeme',
           log,
-        })
-      ).rejects.toThrow('Unable to create Cloud session, token is missing.');
+        },
+        {
+          attemptsCount: 3,
+          attemptDelay: 100,
+        }
+      );
+
+      expect(sessionToken).toBe('mocked_token');
+      expect(axiosRequestMock).toBeCalledTimes(3);
+    });
+
+    test('retries and throws error when response code is not 200', async () => {
+      axiosRequestMock.mockImplementation((config: AxiosRequestConfig) => {
+        if (config.url?.endsWith('/api/v1/saas/auth/_login')) {
+          return Promise.resolve({ data: { message: 'no token' }, status: 503 });
+        }
+        return Promise.reject(new Error(`Unexpected URL: ${config.url}`));
+      });
+
+      await expect(
+        createCloudSession(
+          {
+            hostname: 'cloud',
+            email: 'viewer@elastic.co',
+            password: 'changeme',
+            log,
+          },
+          {
+            attemptsCount: 2,
+            attemptDelay: 100,
+          }
+        )
+      ).rejects.toThrow(
+        `Failed to create the new cloud session: 'POST https://cloud/api/v1/saas/auth/_login' returned 503`
+      );
+      expect(axiosRequestMock).toBeCalledTimes(2);
+    });
+
+    test('retries and throws error when response has no token value', async () => {
+      axiosRequestMock.mockImplementation((config: AxiosRequestConfig) => {
+        if (config.url?.endsWith('/api/v1/saas/auth/_login')) {
+          return Promise.resolve({
+            data: { user_id: 1234, okta_session_id: 5678, authenticated: false },
+            status: 200,
+          });
+        }
+        return Promise.reject(new Error(`Unexpected URL: ${config.url}`));
+      });
+
+      await expect(
+        createCloudSession(
+          {
+            hostname: 'cloud',
+            email: 'viewer@elastic.co',
+            password: 'changeme',
+            log,
+          },
+          {
+            attemptsCount: 3,
+            attemptDelay: 100,
+          }
+        )
+      ).rejects.toThrow(
+        `Failed to create the new cloud session: token is missing in response data\n{"user_id":"REDACTED","okta_session_id":"REDACTED","authenticated":false}`
+      );
+      expect(axiosRequestMock).toBeCalledTimes(3);
+    });
+
+    test(`throws error when retry 'attemptsCount' is below 1`, async () => {
+      await expect(
+        createCloudSession(
+          {
+            hostname: 'cloud',
+            email: 'viewer@elastic.co',
+            password: 'changeme',
+            log,
+          },
+          {
+            attemptsCount: 0,
+            attemptDelay: 100,
+          }
+        )
+      ).rejects.toThrow(
+        'Failed to create the new cloud session, check retry arguments: {"attemptsCount":0,"attemptDelay":100}'
+      );
     });
   });
 
   describe('createSAMLRequest', () => {
+    afterEach(() => {
+      axiosRequestMock.mockClear();
+    });
     test('returns { location, sid }', async () => {
       mockRequestOnce('/internal/security/login', {
         data: {
@@ -130,6 +234,9 @@ describe('saml_auth', () => {
   });
 
   describe('createSAMLResponse', () => {
+    afterEach(() => {
+      axiosGetMock.mockClear();
+    });
     const location = 'https://cloud.test/saml?SAMLRequest=fVLLbtswEPwVgXe9K6%2F';
     const createSAMLResponseParams = {
       location,
@@ -163,6 +270,9 @@ https://kbn.test.co in the same window.`);
   });
 
   describe('finishSAMLHandshake', () => {
+    afterEach(() => {
+      axiosRequestMock.mockClear();
+    });
     const cookieStr = 'mocked_cookie';
     test('returns valid cookie', async () => {
       mockRequestOnce('/api/security/saml/callback', {

packages/kbn-test/src/auth/saml_auth.ts

@@ -17,6 +17,7 @@ import {
   CloudSamlSessionParams,
   CreateSamlSessionParams,
   LocalSamlSessionParams,
+  RetryParams,
   SAMLResponseValueParams,
   UserProfile,
 } from './types';
@@ -34,6 +35,8 @@ export class Session {
   }
 }
 
+const REQUEST_TIMEOUT_MS = 60_000;
+
 const cleanException = (url: string, ex: any) => {
   if (ex.isAxiosError) {
     ex.url = url;
@@ -81,7 +84,13 @@ const getCloudUrl = (hostname: string, pathname: string) => {
   });
 };
 
-export const createCloudSession = async (params: CreateSamlSessionParams) => {
+export const createCloudSession = async (
+  params: CreateSamlSessionParams,
+  retryParams: RetryParams = {
+    attemptsCount: 3,
+    attemptDelay: 15_000,
+  }
+): Promise<string> => {
   const { hostname, email, password, log } = params;
   const cloudLoginUrl = getCloudUrl(hostname, '/api/v1/saas/auth/_login');
   let sessionResponse: AxiosResponse | undefined;
@@ -89,6 +98,7 @@ export const createCloudSession = async (params: CreateSamlSessionParams) => {
     return {
       url: cloudUrl,
       method: 'post',
+      timeout: REQUEST_TIMEOUT_MS,
       data: {
         email,
         password,
@@ -102,24 +112,55 @@ export const createCloudSession = async (params: CreateSamlSessionParams) => {
     };
   };
 
-  try {
-    sessionResponse = await axios.request(requestConfig(cloudLoginUrl));
-  } catch (ex) {
-    log.error(`Failed to create the new cloud session with 'POST ${cloudLoginUrl}'`);
-    cleanException(cloudLoginUrl, ex);
-    throw ex;
+  let attemptsLeft = retryParams.attemptsCount;
+  while (attemptsLeft > 0) {
+    try {
+      sessionResponse = await axios.request(requestConfig(cloudLoginUrl));
+      if (sessionResponse?.status !== 200) {
+        throw new Error(
+          `Failed to create the new cloud session: 'POST ${cloudLoginUrl}' returned ${sessionResponse?.status}`
+        );
+      } else {
+        const token = sessionResponse?.data?.token as string;
+        if (token) {
+          return token;
+        } else {
+          const keysToRedact = ['user_id', 'okta_session_id'];
+          const data = sessionResponse?.data;
+          if (data !== null && typeof data === 'object') {
+            Object.keys(data).forEach((key) => {
+              if (keysToRedact.includes(key)) {
+                data[key] = 'REDACTED';
+              }
+            });
+          }
+          throw new Error(
+            `Failed to create the new cloud session: token is missing in response data\n${JSON.stringify(
+              data
+            )}`
+          );
+        }
+      }
+    } catch (ex) {
+      cleanException(cloudLoginUrl, ex);
+      if (--attemptsLeft > 0) {
+        // log only error message
+        log.error(`${ex.message}\nWaiting ${retryParams.attemptDelay} ms before the next attempt`);
+        await new Promise((resolve) => setTimeout(resolve, retryParams.attemptDelay));
+      } else {
+        log.error(
+          `Failed to create the new cloud session with ${retryParams.attemptsCount} attempts`
+        );
+        // throw original error with stacktrace
+        throw ex;
+      }
+    }
   }
 
-  const token = sessionResponse?.data?.token as string;
-  if (!token) {
-    log.error(
-      `Failed to create cloud session, token is missing in response data: ${JSON.stringify(
-        sessionResponse?.data
-      )}`
-    );
-    throw new Error(`Unable to create Cloud session, token is missing.`);
-  }
-  return token;
+  // should never be reached
+  throw new Error(
+    `Failed to create the new cloud session, check retry arguments: ${JSON.stringify(retryParams)}`
+  );
 };
 
 export const createSAMLRequest = async (kbnUrl: string, kbnVersion: string, log: ToolingLog) => {

packages/kbn-test/src/auth/types.ts

@@ -55,3 +55,8 @@ export interface UserProfile {
   enabled: boolean;
   elastic_cloud_user: boolean;
 }
+
+export interface RetryParams {
+  attemptsCount: number;
+  attemptDelay: number;
+}