[Obs AI Assistant] Add retry statements as an attempt to resolve flak…

…y tests (#192222)
elastic · Nov 13, 2024 · 7323649 · 7323649
1 parent e967193
commit 7323649
Showing 1 changed file with 56 additions and 48 deletions.
diff --git a/...i_assistant_api_integration/tests/knowledge_base/knowledge_base_user_instructions.spec.ts b/...i_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"',
+              },
+            ])
+          );
+        });
       });
     });