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 d436e03554951..3e8cb3123b08b 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 @@ -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 () => {