[Security Solution] improve endpoint metadata tests (elastic#125883)

Co-authored-by: Kibana Machine <42973632+kibanamachine@users.noreply.github.com>
joeypoon · May 19, 2022 · 1ea3fc6 · 1ea3fc6
1 parent efd30bc
commit 1ea3fc6
Show file tree

Hide file tree

Showing 4 changed files with 80 additions and 48 deletions.
diff --git a/x-pack/plugins/security_solution/common/endpoint/data_loaders/index_fleet_agent.ts b/x-pack/plugins/security_solution/common/endpoint/data_loaders/index_fleet_agent.ts
@@ -23,7 +23,7 @@ import { wrapErrorAndRejectPromise } from './utils';
 const defaultFleetAgentGenerator = new FleetAgentGenerator();
 
 export interface IndexedFleetAgentResponse {
-  agents: Agent[];
+  agents: Array<Agent & FleetServerAgent>;
   fleetAgentsIndex: string;
 }
 

diff --git a/x-pack/test/security_solution_endpoint/services/endpoint.ts b/x-pack/test/security_solution_endpoint/services/endpoint.ts
@@ -11,6 +11,7 @@ import {
   metadataCurrentIndexPattern,
   metadataTransformPrefix,
   METADATA_UNITED_INDEX,
+  METADATA_UNITED_TRANSFORM,
 } from '@kbn/security-solution-plugin/common/endpoint/constants';
 import {
   deleteIndexedHostsAndAlerts,
@@ -77,6 +78,27 @@ export class EndpointTestResources extends FtrService {
     await this.transform.api.updateTransform(transform.id, { frequency }).catch(catchAndWrapError);
   }
 
+  private async stopTransform(transformId: string) {
+    const stopRequest = {
+      transform_id: `${transformId}*`,
+      force: true,
+      wait_for_completion: true,
+      allow_no_match: true,
+    };
+    return this.esClient.transform.stopTransform(stopRequest);
+  }
+
+  private async startTransform(transformId: string) {
+    const transformsResponse = await this.esClient.transform.getTransform({
+      transform_id: `${transformId}*`,
+    });
+    return Promise.all(
+      transformsResponse.transforms.map((transform) => {
+        return this.esClient.transform.startTransform({ transform_id: transform.id });
+      })
+    );
+  }
+
   /**
    * Loads endpoint host/alert/event data into elasticsearch
    * @param [options]
@@ -86,6 +108,8 @@ export class EndpointTestResources extends FtrService {
    * @param [options.enableFleetIntegration=true] When set to `true`, Fleet data will also be loaded (ex. Integration Policies, Agent Policies, "fake" Agents)
    * @param [options.generatorSeed='seed`] The seed to be used by the data generator. Important in order to ensure the same data is generated on very run.
    * @param [options.waitUntilTransformed=true] If set to `true`, the data loading process will wait until the endpoint hosts metadata is processed by the transform
+   * @param [options.waitTimeout=60000] If waitUntilTransformed=true, number of ms to wait until timeout
+   * @param [options.customIndexFn] If provided, will use this function to generate and index data instead
    */
   async loadEndpointData(
     options: Partial<{
@@ -95,6 +119,8 @@ export class EndpointTestResources extends FtrService {
       enableFleetIntegration: boolean;
       generatorSeed: string;
       waitUntilTransformed: boolean;
+      waitTimeout: number;
+      customIndexFn: () => Promise<IndexedHostsAndAlertsResponse>;
     }> = {}
   ): Promise<IndexedHostsAndAlertsResponse> {
     const {
@@ -104,25 +130,39 @@ export class EndpointTestResources extends FtrService {
       enableFleetIntegration = true,
       generatorSeed = 'seed',
       waitUntilTransformed = true,
+      waitTimeout = 60000,
+      customIndexFn,
     } = options;
 
+    if (waitUntilTransformed) {
+      // need this before indexing docs so that the united transform doesn't
+      // create a checkpoint with a timestamp after the doc timestamps
+      await this.stopTransform(METADATA_UNITED_TRANSFORM);
+    }
+
     // load data into the system
-    const indexedData = await indexHostsAndAlerts(
-      this.esClient as Client,
-      this.kbnClient,
-      generatorSeed,
-      numHosts,
-      numHostDocs,
-      'metrics-endpoint.metadata-default',
-      'metrics-endpoint.policy-default',
-      'logs-endpoint.events.process-default',
-      'logs-endpoint.alerts-default',
-      alertsPerHost,
-      enableFleetIntegration
-    );
+    const indexedData = customIndexFn
+      ? await customIndexFn()
+      : await indexHostsAndAlerts(
+          this.esClient as Client,
+          this.kbnClient,
+          generatorSeed,
+          numHosts,
+          numHostDocs,
+          'metrics-endpoint.metadata-default',
+          'metrics-endpoint.policy-default',
+          'logs-endpoint.events.process-default',
+          'logs-endpoint.alerts-default',
+          alertsPerHost,
+          enableFleetIntegration
+        );
 
     if (waitUntilTransformed) {
-      await this.waitForEndpoints(indexedData.hosts.map((host) => host.agent.id));
+      const metadataIds = Array.from(new Set(indexedData.hosts.map((host) => host.agent.id)));
+      await this.waitForEndpoints(metadataIds, waitTimeout);
+      await this.startTransform(METADATA_UNITED_TRANSFORM);
+      const agentIds = Array.from(new Set(indexedData.agents.map((agent) => agent.agent!.id)));
+      await this.waitForUnitedEndpoints(agentIds, waitTimeout);
     }
 
     return indexedData;

diff --git a/x-pack/test/security_solution_endpoint_api_int/apis/endpoint_authz.ts b/x-pack/test/security_solution_endpoint_api_int/apis/endpoint_authz.ts
@@ -5,7 +5,6 @@
  * 2.0.
  */
 
-import { IndexedHostsAndAlertsResponse } from '@kbn/security-solution-plugin/common/endpoint/index_data';
 import { wrapErrorAndRejectPromise } from '@kbn/security-solution-plugin/common/endpoint/data_loaders/utils';
 import { FtrProviderContext } from '../ftr_provider_context';
 import {
@@ -15,23 +14,15 @@ import {
 } from '../../common/services/security_solution';
 
 export default function ({ getService }: FtrProviderContext) {
-  const endpointTestResources = getService('endpointTestResources');
   const supertestWithoutAuth = getService('supertestWithoutAuth');
 
   describe('When attempting to call an endpoint api with no authz', () => {
-    let loadedData: IndexedHostsAndAlertsResponse;
-
     before(async () => {
       // create role/user
       await createUserAndRole(getService, ROLES.t1_analyst);
-      loadedData = await endpointTestResources.loadEndpointData();
     });
 
     after(async () => {
-      if (loadedData) {
-        await endpointTestResources.unloadEndpointData(loadedData);
-      }
-
       // delete role/user
       await deleteUserAndRole(getService, ROLES.t1_analyst);
     });

diff --git a/x-pack/test/security_solution_endpoint_api_int/apis/metadata.ts b/x-pack/test/security_solution_endpoint_api_int/apis/metadata.ts
@@ -19,6 +19,8 @@ import {
 import { AGENTS_INDEX } from '@kbn/fleet-plugin/common';
 import { indexFleetEndpointPolicy } from '@kbn/security-solution-plugin/common/endpoint/data_loaders/index_fleet_endpoint_policy';
 import { TRANSFORM_STATES } from '@kbn/security-solution-plugin/common/constants';
+import type { IndexedHostsAndAlertsResponse } from '@kbn/security-solution-plugin/common/endpoint/index_data';
+
 import { generateAgentDocs, generateMetadataDocs } from './metadata.fixtures';
 import {
   deleteAllDocsFromMetadataCurrentIndex,
@@ -47,38 +49,37 @@ export default function ({ getService }: FtrProviderContext) {
         const numberOfHostsInFixture = 2;
 
         before(async () => {
-          await stopTransform(getService, `${METADATA_UNITED_TRANSFORM}*`);
           await deleteAllDocsFromFleetAgents(getService);
           await deleteAllDocsFromMetadataDatastream(getService);
           await deleteAllDocsFromMetadataCurrentIndex(getService);
           await deleteAllDocsFromIndex(getService, METADATA_UNITED_INDEX);
 
-          // generate an endpoint policy and attach id to agents since
-          // metadata list api filters down to endpoint policies only
-          const policy = await indexFleetEndpointPolicy(
-            getService('kibanaServer'),
-            `Default ${uuid.v4()}`,
-            '1.1.1'
-          );
-          const policyId = policy.integrationPolicies[0].policy_id;
-          const currentTime = new Date().getTime();
+          const customIndexFn = async (): Promise<IndexedHostsAndAlertsResponse> => {
+            // generate an endpoint policy and attach id to agents since
+            // metadata list api filters down to endpoint policies only
+            const policy = await indexFleetEndpointPolicy(
+              getService('kibanaServer'),
+              `Default ${uuid.v4()}`,
+              '1.1.1'
+            );
+            const policyId = policy.integrationPolicies[0].policy_id;
+            const currentTime = new Date().getTime();
 
-          const agentDocs = generateAgentDocs(currentTime, policyId);
+            const agentDocs = generateAgentDocs(currentTime, policyId);
+            const metadataDocs = generateMetadataDocs(currentTime);
 
-          await Promise.all([
-            bulkIndex(getService, AGENTS_INDEX, agentDocs),
-            bulkIndex(getService, METADATA_DATASTREAM, generateMetadataDocs(currentTime)),
-          ]);
+            await Promise.all([
+              bulkIndex(getService, AGENTS_INDEX, agentDocs),
+              bulkIndex(getService, METADATA_DATASTREAM, metadataDocs),
+            ]);
 
-          await endpointTestResources.waitForEndpoints(
-            agentDocs.map((doc) => doc.agent.id),
-            60000
-          );
-          await startTransform(getService, METADATA_UNITED_TRANSFORM);
-          await endpointTestResources.waitForUnitedEndpoints(
-            agentDocs.map((doc) => doc.agent.id),
-            60000
-          );
+            return {
+              agents: agentDocs,
+              hosts: metadataDocs,
+            } as unknown as IndexedHostsAndAlertsResponse;
+          };
+
+          await endpointTestResources.loadEndpointData({ customIndexFn });
         });
 
         after(async () => {