From 328871829c6dd4eea4d0e42ec82bce8ba8828896 Mon Sep 17 00:00:00 2001
From: Viduni Wickramarachchi <viduni.wickramarachchi@elastic.co>
Date: Thu, 14 Nov 2024 11:09:42 -0500
Subject: [PATCH] [Obs AI Assistant] Add retry statements as an attempt to
 resolve flaky tests (#200022)

Closes https://github.com/elastic/kibana/issues/192222

## Summary

### Problem
The test appears to be flaky, potentially because the entries are not
available at the time of retrieval. This cannot be reproduced locally or
via the flaky test runner. (more details
[here](https://github.com/elastic/kibana/pull/196026#issuecomment-2409056856))

### Solution
Add a retry when fetching the instructions and check whether the number
of instructions returned by the API endpoint is the same number of
instructions expected.

### Checklist

- [x] [Flaky Test
Runner](https://ci-stats.kibana.dev/trigger_flaky_test_runner/1) was
used on any tests changed

(cherry picked from commit 53c05a33e7fe917b7d64b7d3ca664aaf3f9cf3ca)
---
 .../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 a9c1f245a1ac3..58c6b4d38be4f 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"',
+              },
+            ])
+          );
+        });
       });
     });