[Security GenAI] Fix and un-skip Knowledge Base Integration Tests

x-pack/plugins/elastic_assistant/server/routes/knowledge_base/get_knowledge_base_status.ts

@@ -61,12 +61,11 @@ export const getKnowledgeBaseStatusRoute = (router: ElasticAssistantPluginRouter
 
           const indexExists = true; // Installed at startup, always true
           const pipelineExists = true; // Installed at startup, always true
-          const modelExists = await kbDataClient.isModelInstalled();
           const setupAvailable = await kbDataClient.isSetupAvailable();
           const isModelDeployed = await kbDataClient.isModelDeployed();
 
           const body: ReadKnowledgeBaseResponse = {
-            elser_exists: modelExists,
+            elser_exists: isModelDeployed,
             index_exists: indexExists,
             is_setup_in_progress: kbDataClient.isSetupInProgress,
             is_setup_available: setupAvailable,

x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/trial_license_complete_tier/entries.ts

@@ -7,6 +7,7 @@
 
 import expect from 'expect';
 import { KNOWLEDGE_BASE_ENTRIES_TABLE_MAX_PAGE_SIZE } from '@kbn/elastic-assistant-plugin/common/constants';
+import { isSystemEntry } from '@kbn/elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/helpers';
 import { FtrProviderContext } from '../../../../../ftr_provider_context';
 import { createEntry, createEntryForUser } from '../utils/create_entry';
 import { findEntries } from '../utils/find_entry';
@@ -15,6 +16,7 @@ import {
   deleteTinyElser,
   installTinyElser,
   setupKnowledgeBase,
+  waitForKnowledgeBaseModelToBeDeployed,
 } from '../utils/helpers';
 import { removeServerGeneratedProperties } from '../utils/remove_server_generated_properties';
 import { MachineLearningProvider } from '../../../../../../functional/services/ml';
@@ -32,10 +34,11 @@ export default ({ getService }: FtrProviderContext) => {
   const es = getService('es');
   const ml = getService('ml') as ReturnType<typeof MachineLearningProvider>;
 
-  describe.skip('@ess Basic Security AI Assistant Knowledge Base Entries', () => {
+  describe('@ess Basic Security AI Assistant Knowledge Base Entries', () => {
     before(async () => {
       await installTinyElser(ml);
-      await setupKnowledgeBase(supertest, log);
+      setupKnowledgeBase(supertest, log);
+      await waitForKnowledgeBaseModelToBeDeployed(supertest, log);
     });
 
     after(async () => {
@@ -170,8 +173,9 @@ export default ({ getService }: FtrProviderContext) => {
         );
 
         const entries = await findEntries({ supertest, log });
+        const globalEntries = entries.data.filter((entry) => !isSystemEntry(entry));
 
-        expect(entries.total).toEqual(1);
+        expect(globalEntries.length).toEqual(1);
       });
 
       it('should not see other users private entries', async () => {
@@ -189,8 +193,9 @@ export default ({ getService }: FtrProviderContext) => {
         );
 
         const entries = await findEntries({ supertest, log });
+        const privateEntries = entries.data.filter((entry) => !isSystemEntry(entry));
 
-        expect(entries.total).toEqual(0);
+        expect(privateEntries.length).toEqual(0);
       });
     });
 

x-pack/test/security_solution_api_integration/test_suites/genai/knowledge_base/entries/utils/helpers.ts

@@ -9,6 +9,7 @@ import { Client } from '@elastic/elasticsearch';
 import {
   CreateKnowledgeBaseResponse,
   ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL,
+  ReadKnowledgeBaseResponse,
 } from '@kbn/elastic-assistant-common';
 
 import { ELASTIC_HTTP_VERSION_HEADER } from '@kbn/core-http-common';
@@ -17,7 +18,7 @@ import type SuperTest from 'supertest';
 import { MachineLearningProvider } from '../../../../../../functional/services/ml';
 import { SUPPORTED_TRAINED_MODELS } from '../../../../../../functional/services/ml/api';
 
-import { routeWithNamespace } from '../../../../../../common/utils/security_solution';
+import { routeWithNamespace, waitFor } from '../../../../../../common/utils/security_solution';
 
 export const TINY_ELSER = {
   ...SUPPORTED_TRAINED_MODELS.TINY_ELSER,
@@ -81,6 +82,56 @@ export const setupKnowledgeBase = async (
   }
 };
 
+/**
+ * Get Knowledge Base status
+ * @param supertest The supertest deps
+ * @param log The tooling logger
+ * @param resource
+ * @param namespace The Kibana Space where the KB should be set up
+ */
+export const getKnowledgeBaseStatus = async (
+  supertest: SuperTest.Agent,
+  log: ToolingLog,
+  resource?: string,
+  namespace?: string
+): Promise<ReadKnowledgeBaseResponse> => {
+  const path = ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL.replace('{resource?}', resource || '');
+  const route = routeWithNamespace(path, namespace);
+  const response = await supertest
+    .get(route)
+    .set('kbn-xsrf', 'true')
+    .set(ELASTIC_HTTP_VERSION_HEADER, '1')
+    .send();
+  if (response.status !== 200) {
+    throw new Error(
+      `Unexpected non 200 ok when attempting to get Knowledge Base status: ${JSON.stringify(
+        response.status
+      )},${JSON.stringify(response, null, 4)}`
+    );
+  } else {
+    return response.body;
+  }
+};
+
+/**
+ * Waits for the knowledge base model to be deployed
+ * @param supertest The supertest deps
+ * @param log The tooling logger
+ */
+export const waitForKnowledgeBaseModelToBeDeployed = async (
+  supertest: SuperTest.Agent,
+  log: ToolingLog
+): Promise<void> => {
+  await waitFor(
+    async () => {
+      const status = await getKnowledgeBaseStatus(supertest, log);
+      return !!status.elser_exists;
+    },
+    'waitForKnowledgeBaseModelToBeDeployed',
+    log
+  );
+};
+
 /**
  * Clear Knowledge Base
  * @param es

x-pack/test/security_solution_api_integration/tsconfig.json

@@ -52,5 +52,6 @@
     "@kbn/ftr-common-functional-ui-services",
     "@kbn/spaces-plugin",
     "@kbn/elastic-assistant-plugin",
+    "@kbn/elastic-assistant",
   ]
 }