[Obs AI Assistant] Attempt to fix flaky KB migration test (elastic#20…

…2926)
viduni94 · Jan 13, 2025 · e3c7e43 · e3c7e43
1 parent 8bc22a1
commit e3c7e43
Showing 1 changed file with 29 additions and 26 deletions.
diff --git a/..._agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_migration.spec.ts b/..._agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_migration.spec.ts
@@ -24,6 +24,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
   const esArchiver = getService('esArchiver');
   const es = getService('es');
   const ml = getService('ml');
+  const retry = getService('retry');
 
   const archive =
     'x-pack/test/functional/es_archives/observability/ai_assistant/knowledge_base_8_15';
@@ -91,32 +92,34 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
       });
 
       it('the docs have semantic_text embeddings', async () => {
-        const hits = await getKnowledgeBaseEntries();
-        const hasSemanticTextEmbeddings = hits.every((hit) => hit._source?.semantic_text);
-        expect(hasSemanticTextEmbeddings).to.be(true);
-
-        expect(
-          orderBy(hits, '_source.title').map(({ _source }) => {
-            const { text, inference } = _source?.semantic_text!;
-
-            return {
-              text,
-              inferenceId: inference.inference_id,
-              chunkCount: inference.chunks.length,
-            };
-          })
-        ).to.eql([
-          {
-            text: 'To infinity and beyond!',
-            inferenceId: AI_ASSISTANT_KB_INFERENCE_ID,
-            chunkCount: 1,
-          },
-          {
-            text: "The user's favourite color is blue.",
-            inferenceId: AI_ASSISTANT_KB_INFERENCE_ID,
-            chunkCount: 1,
-          },
-        ]);
+        await retry.try(async () => {
+          const hits = await getKnowledgeBaseEntries();
+          const hasSemanticTextEmbeddings = hits.every((hit) => hit._source?.semantic_text);
+          expect(hasSemanticTextEmbeddings).to.be(true);
+
+          expect(
+            orderBy(hits, '_source.title').map(({ _source }) => {
+              const { text, inference } = _source?.semantic_text!;
+
+              return {
+                text,
+                inferenceId: inference.inference_id,
+                chunkCount: inference.chunks.length,
+              };
+            })
+          ).to.eql([
+            {
+              text: 'To infinity and beyond!',
+              inferenceId: AI_ASSISTANT_KB_INFERENCE_ID,
+              chunkCount: 1,
+            },
+            {
+              text: "The user's favourite color is blue.",
+              inferenceId: AI_ASSISTANT_KB_INFERENCE_ID,
+              chunkCount: 1,
+            },
+          ]);
+        });
       });
 
       it('returns entries correctly via API', async () => {