[Obs AI Assistant] Retry creating entries if any entry is missing (#1…

…92222)
elastic · Oct 25, 2024 · 9552435 · 9552435
1 parent 479ebea
commit 9552435
Showing 1 changed file with 64 additions and 16 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
@@ -29,6 +29,7 @@ export default function ApiTest({ getService }: FtrProviderContext) {
   const es = getService('es');
   const ml = getService('ml');
   const log = getService('log');
+  const retry = getService('retry');
 
   describe('Knowledge base user instructions', () => {
     const userJohn = 'john';
@@ -54,41 +55,88 @@ export default function ApiTest({ getService }: FtrProviderContext) {
     describe('when creating private and public user instructions', () => {
       before(async () => {
         await clearKnowledgeBase(es);
-        await new Promise((resolve) => setTimeout(resolve, 500));
 
-        const promises = [
+        const knowledgeBaseEntries = [
           {
+            doc_id: 'public-doc-from-editor',
+            text: 'Public user instruction from "editor"',
             username: 'editor',
             isPublic: true,
           },
           {
+            doc_id: 'private-doc-from-editor',
+            text: 'Private user instruction from "editor"',
             username: 'editor',
             isPublic: false,
           },
           {
+            doc_id: 'public-doc-from-john',
+            text: 'Public user instruction from "john"',
             username: userJohn,
             isPublic: true,
           },
           {
+            doc_id: 'private-doc-from-john',
+            text: 'Private user instruction from "john"',
             username: userJohn,
             isPublic: false,
           },
-        ].map(async ({ username, isPublic }) => {
-          const visibility = isPublic ? 'Public' : 'Private';
-          await getScopedApiClientForUsername(username)({
-            endpoint: 'PUT /internal/observability_ai_assistant/kb/user_instructions',
-            params: {
-              body: {
-                id: `${visibility.toLowerCase()}-doc-from-${username}`,
-                text: `${visibility} user instruction from "${username}"`,
-                public: isPublic,
+        ];
+
+        // Initial attempt to create all entries
+        await Promise.all(
+          knowledgeBaseEntries.map(async ({ doc_id: id, text, username, isPublic }) => {
+            await getScopedApiClientForUsername(username)({
+              endpoint: 'PUT /internal/observability_ai_assistant/kb/user_instructions',
+              params: {
+                body: {
+                  id,
+                  text,
+                  public: isPublic,
+                },
               },
-            },
-          }).expect(200);
-        });
+            }).expect(200);
+          })
+        );
 
-        await Promise.all(promises);
-        await new Promise((resolve) => setTimeout(resolve, 500));
+        // To resolve flakiness, retry the logic to only add the missing entries
+        await retry.try(async () => {
+          const editorEntries = await observabilityAIAssistantAPIClient.editorUser({
+            endpoint: 'GET /internal/observability_ai_assistant/kb/user_instructions',
+          });
+          const userJohnEntries = await getScopedApiClientForUsername(userJohn)({
+            endpoint: 'GET /internal/observability_ai_assistant/kb/user_instructions',
+          });
+
+          const existingIds = [
+            ...editorEntries.body.userInstructions.map((entry) => entry.doc_id),
+            ...userJohnEntries.body.userInstructions.map((entry) => entry.doc_id),
+          ];
+
+          const missingEntries = knowledgeBaseEntries.filter(
+            (entry) => !existingIds.includes(entry.doc_id)
+          );
+
+          if (missingEntries.length === 0) {
+            return; // All entries are present, exit retry
+          }
+
+          // Retry creating only the missing entries
+          await Promise.all(
+            missingEntries.map(async ({ doc_id: id, text, username, isPublic }) => {
+              await getScopedApiClientForUsername(username)({
+                endpoint: 'PUT /internal/observability_ai_assistant/kb/user_instructions',
+                params: {
+                  body: {
+                    id,
+                    text,
+                    public: isPublic,
+                  },
+                },
+              }).expect(200);
+            })
+          );
+        });
       });
 
       it('"editor" can retrieve their own private instructions and the public instruction', async () => {