From 7323649a78a7b59d6a22968b51d5903474456cc3 Mon Sep 17 00:00:00 2001
From: Viduni Wickramarachchi <viduni.wickramarachchi@elastic.co>
Date: Wed, 13 Nov 2024 10:05:41 -0500
Subject: [PATCH] [Obs AI Assistant] Add retry statements as an attempt to
 resolve flaky tests (#192222)

---
 .../knowledge_base_user_instructions.spec.ts  | 104 ++++++++++--------
 1 file changed, 56 insertions(+), 48 deletions(-)

diff --git a/x-pack/test/observability_ai_assistant_api_integration/tests/knowledge_base/knowledge_base_user_instructions.spec.ts b/x-pack/test/observability_ai_assistant_api_integration/tests/knowledge_base/knowledge_base_user_instructions.spec.ts
index a9c1f245a1ac..58c6b4d38be4 100644
--- a/x-pack/test/observability_ai_assistant_api_integration/tests/knowledge_base/knowledge_base_user_instructions.spec.ts
+++ b/x-pack/test/observability_ai_assistant_api_integration/tests/knowledge_base/knowledge_base_user_instructions.spec.ts
@@ -28,6 +28,7 @@ export default function ApiTest({ getService }: FtrProviderContext) {
   const es = getService('es');
   const ml = getService('ml');
   const log = getService('log');
+  const retry = getService('retry');
   const getScopedApiClientForUsername = getService('getScopedApiClientForUsername');
 
   describe('Knowledge base user instructions', () => {
@@ -85,62 +86,69 @@ export default function ApiTest({ getService }: FtrProviderContext) {
       });
 
       it('"editor" can retrieve their own private instructions and the public instruction', async () => {
-        const res = await observabilityAIAssistantAPIClient.editor({
-          endpoint: 'GET /internal/observability_ai_assistant/kb/user_instructions',
-        });
+        await retry.try(async () => {
+          const res = await observabilityAIAssistantAPIClient.editor({
+            endpoint: 'GET /internal/observability_ai_assistant/kb/user_instructions',
+          });
 
-        const instructions = res.body.userInstructions;
+          const instructions = res.body.userInstructions;
+          expect(instructions).to.have.length(3);
 
-        const sortById = (data: Array<Instruction & { public?: boolean }>) => sortBy(data, 'id');
+          const sortById = (data: Array<Instruction & { public?: boolean }>) => sortBy(data, 'id');
 
-        expect(sortById(instructions)).to.eql(
-          sortById([
-            {
-              id: 'private-doc-from-editor',
-              public: false,
-              text: 'Private user instruction from "editor"',
-            },
-            {
-              id: 'public-doc-from-editor',
-              public: true,
-              text: 'Public user instruction from "editor"',
-            },
-            {
-              id: 'public-doc-from-secondary_editor',
-              public: true,
-              text: 'Public user instruction from "secondary_editor"',
-            },
-          ])
-        );
+          expect(sortById(instructions)).to.eql(
+            sortById([
+              {
+                id: 'private-doc-from-editor',
+                public: false,
+                text: 'Private user instruction from "editor"',
+              },
+              {
+                id: 'public-doc-from-editor',
+                public: true,
+                text: 'Public user instruction from "editor"',
+              },
+              {
+                id: 'public-doc-from-secondary_editor',
+                public: true,
+                text: 'Public user instruction from "secondary_editor"',
+              },
+            ])
+          );
+        });
       });
 
       it('"secondaryEditor" can retrieve their own private instructions and the public instruction', async () => {
-        const res = await observabilityAIAssistantAPIClient.secondaryEditor({
-          endpoint: 'GET /internal/observability_ai_assistant/kb/user_instructions',
-        });
-        const instructions = res.body.userInstructions;
+        await retry.try(async () => {
+          const res = await observabilityAIAssistantAPIClient.secondaryEditor({
+            endpoint: 'GET /internal/observability_ai_assistant/kb/user_instructions',
+          });
 
-        const sortById = (data: Array<Instruction & { public?: boolean }>) => sortBy(data, 'id');
+          const instructions = res.body.userInstructions;
+          expect(instructions).to.have.length(3);
 
-        expect(sortById(instructions)).to.eql(
-          sortById([
-            {
-              id: 'public-doc-from-editor',
-              public: true,
-              text: 'Public user instruction from "editor"',
-            },
-            {
-              id: 'public-doc-from-secondary_editor',
-              public: true,
-              text: 'Public user instruction from "secondary_editor"',
-            },
-            {
-              id: 'private-doc-from-secondary_editor',
-              public: false,
-              text: 'Private user instruction from "secondary_editor"',
-            },
-          ])
-        );
+          const sortById = (data: Array<Instruction & { public?: boolean }>) => sortBy(data, 'id');
+
+          expect(sortById(instructions)).to.eql(
+            sortById([
+              {
+                id: 'public-doc-from-editor',
+                public: true,
+                text: 'Public user instruction from "editor"',
+              },
+              {
+                id: 'public-doc-from-secondary_editor',
+                public: true,
+                text: 'Public user instruction from "secondary_editor"',
+              },
+              {
+                id: 'private-doc-from-secondary_editor',
+                public: false,
+                text: 'Private user instruction from "secondary_editor"',
+              },
+            ])
+          );
+        });
       });
     });