tests

x-pack/plugins/security_solution_serverless/server/cloud_security/cloud_security_metering_task.test.ts

@@ -16,7 +16,8 @@ import {
 import type { ServerlessSecurityConfig } from '../config';
 import type { CloudSecuritySolutions } from './types';
 import type { ProductTier } from '../../common/product';
-import { CLOUD_SECURITY_TASK_TYPE, CSPM, KSPM, CNVM } from './constants';
+import { CLOUD_SECURITY_TASK_TYPE, CSPM, KSPM, CNVM, CLOUD_DEFEND } from './constants';
+import { getCloudDefendUsageRecords } from './defend_for_containers_metering';
 
 const mockEsClient = elasticsearchServiceMock.createStart().client.asInternalUser;
 const logger: ReturnType<typeof loggingSystemMock.createLogger> = loggingSystemMock.createLogger();
@@ -160,26 +161,6 @@ describe('getCloudSecurityUsageRecord', () => {
 });
 
 describe('getSearchQueryByCloudSecuritySolution', () => {
-  it('should return the correct search query for CLOUD_DEFEND', () => {
-    const searchFrom = new Date('2023-07-30T15:11:41.738Z');
-
-    const result = getSearchQueryByCloudSecuritySolution('cloud_defend', searchFrom);
-
-    expect(result).toEqual({
-      bool: {
-        must: [
-          {
-            range: {
-              '@timestamp': {
-                gt: '2023-07-30T15:11:41.738Z',
-              },
-            },
-          },
-        ],
-      },
-    });
-  });
-
   it('should return the correct search query for CSPM', () => {
     const searchFrom = new Date('2023-07-30T15:11:41.738Z');
 
@@ -296,38 +277,42 @@ describe('should return the relevant product tier', () => {
     expect(tier).toBe('complete');
   });
 
-  it('should return usageRecords with correct values for cloud defend', async () => {
+  it('should return none tier in case cloud product line is missing ', async () => {
+    const serverlessSecurityConfig = {
+      enabled: true,
+      developer: {},
+      productTypes: [{ product_line: 'endpoint', product_tier: 'complete' }],
+    } as unknown as ServerlessSecurityConfig;
+
+    const tier = getCloudProductTier(serverlessSecurityConfig, logger);
+
+    expect(tier).toBe('none');
+  });
+});
+
+describe('cloud defend metering', () => {
+  it('should return usageRecords with correct values', async () => {
     const cloudSecuritySolution = 'cloud_defend';
-    // @ts-ignore
-    mockEsClient.search.mockResolvedValueOnce({
-      hits: { hits: [{ _id: 'someRecord', _index: 'mockIndex' }] }, // mocking for indexHasDataInDateRange
-    });
+    const agentId1 = chance.guid();
+    const eventIngestedStr = '2024-05-28T12:10:51Z';
+    const eventIngestedTimestamp = new Date(eventIngestedStr);
 
     // @ts-ignore
     mockEsClient.search.mockResolvedValueOnce({
-      aggregations: {
-        asset_count_groups: {
-          buckets: [
-            {
-              key_as_string: 'true',
-              unique_assets: {
-                value: 10,
-              },
-              min_timestamp: {
-                value_as_string: '2023-07-30T15:11:41.738Z',
-              },
-            },
-            {
-              key_as_string: 'false',
-              unique_assets: {
-                value: 5,
-              },
-              min_timestamp: {
-                value_as_string: '2023-07-30T15:11:41.738Z',
+      hits: {
+        hits: [
+          {
+            _id: 'someRecord',
+            _index: 'mockIndex',
+            _source: {
+              'cloud_defend.block_action_enabled': true,
+              'agent.id': agentId1,
+              event: {
+                ingested: eventIngestedStr,
               },
             },
-          ],
-        },
+          },
+        ],
       },
     });
 
@@ -336,48 +321,33 @@ describe('should return the relevant product tier', () => {
 
     const tier = 'essentials' as ProductTier;
 
-    const result = await getCloudSecurityUsageRecord({
+    const result = await getCloudDefendUsageRecords({
       esClient: mockEsClient,
       projectId,
-      logger,
       taskId,
       lastSuccessfulReport: new Date(),
       cloudSecuritySolution,
+      logger,
       tier,
     });
 
+    const roundedIngestedTimestamp = eventIngestedTimestamp;
+    roundedIngestedTimestamp.setMinutes(0);
+    roundedIngestedTimestamp.setSeconds(0);
+    roundedIngestedTimestamp.setMilliseconds(0);
+
     expect(result).toEqual([
       {
         id: expect.stringContaining(
-          `${CLOUD_SECURITY_TASK_TYPE}_${cloudSecuritySolution}_${projectId}`
+          `${projectId}_${agentId1}_${roundedIngestedTimestamp.toISOString()}`
         ),
-        usage_timestamp: '2023-07-30T15:11:41.738Z',
-        creation_timestamp: expect.any(String), // Expect a valid ISO string
+        usage_timestamp: eventIngestedStr,
+        creation_timestamp: expect.any(String),
         usage: {
           type: CLOUD_SECURITY_TASK_TYPE,
-          sub_type: `${cloudSecuritySolution}_block_action_enabled_true`,
-          quantity: 10,
-          period_seconds: expect.any(Number),
-        },
-        source: {
-          id: taskId,
-          instance_group_id: projectId,
-          metadata: {
-            tier: 'essentials',
-          },
-        },
-      },
-      {
-        id: expect.stringContaining(
-          `${CLOUD_SECURITY_TASK_TYPE}_${cloudSecuritySolution}_${projectId}`
-        ),
-        usage_timestamp: '2023-07-30T15:11:41.738Z',
-        creation_timestamp: expect.any(String), // Expect a valid ISO string
-        usage: {
-          type: CLOUD_SECURITY_TASK_TYPE,
-          sub_type: `${cloudSecuritySolution}_block_action_enabled_false`,
-          quantity: 5,
-          period_seconds: expect.any(Number),
+          sub_type: CLOUD_DEFEND,
+          quantity: 1,
+          period_seconds: 3600,
         },
         source: {
           id: taskId,
@@ -390,15 +360,102 @@ describe('should return the relevant product tier', () => {
     ]);
   });
 
-  it('should return none tier in case cloud product line is missing ', async () => {
-    const serverlessSecurityConfig = {
-      enabled: true,
-      developer: {},
-      productTypes: [{ product_line: 'endpoint', product_tier: 'complete' }],
-    } as unknown as ServerlessSecurityConfig;
+  it('should return an empty array when Elasticsearch returns an empty response', async () => {
+    // @ts-ignore
+    mockEsClient.search.mockResolvedValueOnce({
+      hits: {
+        hits: [],
+      },
+    });
+    const tier = 'essentials' as ProductTier;
+    // Call the function with mock parameters
+    const result = await getCloudDefendUsageRecords({
+      esClient: mockEsClient,
+      projectId: chance.guid(),
+      taskId: chance.guid(),
+      lastSuccessfulReport: new Date(),
+      cloudSecuritySolution: 'cloud_defend',
+      logger,
+      tier,
+    });
 
-    const tier = getCloudProductTier(serverlessSecurityConfig, logger);
+    // Assert that the result is an empty array
+    expect(result).toEqual([]);
+  });
 
-    expect(tier).toBe('none');
+  it('should handle errors from Elasticsearch', async () => {
+    // Mock Elasticsearch client's search method to throw an error
+    mockEsClient.search.mockRejectedValueOnce(new Error('Elasticsearch query failed'));
+
+    const tier = 'essentials' as ProductTier;
+
+    // Call the function with mock parameters
+    await getCloudDefendUsageRecords({
+      esClient: mockEsClient,
+      projectId: chance.guid(),
+      taskId: chance.guid(),
+      lastSuccessfulReport: new Date(),
+      cloudSecuritySolution: 'cloud_defend',
+      logger,
+      tier,
+    });
+
+    // Assert that the logger's error method was called with the correct error message
+    expect(logger.error).toHaveBeenCalledWith(
+      'Failed to fetch cloud_defend metering data Error: Elasticsearch query failed'
+    );
+  });
+
+  it('should return usageRecords when Elasticsearch returns multiple records', async () => {
+    // Mock Elasticsearch response with multiple records
+    const agentId1 = chance.guid();
+    const agentId2 = chance.guid();
+    const eventIngestedStr1 = '2024-05-28T12:10:51Z';
+    const eventIngestedStr2 = '2024-05-28T13:10:51Z';
+
+    // @ts-ignore
+    mockEsClient.search.mockResolvedValueOnce({
+      hits: {
+        hits: [
+          {
+            _id: 'record1',
+            _index: 'mockIndex',
+            _source: {
+              'cloud_defend.block_action_enabled': true,
+              'agent.id': agentId1,
+              event: {
+                ingested: eventIngestedStr1,
+              },
+            },
+          },
+          {
+            _id: 'record2',
+            _index: 'mockIndex',
+            _source: {
+              'cloud_defend.block_action_enabled': true,
+              'agent.id': agentId2,
+              event: {
+                ingested: eventIngestedStr2,
+              },
+            },
+          },
+        ],
+      },
+    });
+    const tier = 'essentials' as ProductTier;
+
+    // Call the function with mock parameters
+    const result = await getCloudDefendUsageRecords({
+      esClient: mockEsClient,
+      projectId: chance.guid(),
+      taskId: chance.guid(),
+      lastSuccessfulReport: new Date(),
+      cloudSecuritySolution: 'cloud_defend',
+      logger,
+      tier,
+    });
+
+    // Assert that the result contains usage records for both records
+    expect(result).toHaveLength(2);
   });
 });

x-pack/plugins/security_solution_serverless/server/cloud_security/defend_for_containers_metering.ts

@@ -6,12 +6,17 @@
  */
 
 import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server';
-import type { AggregationsAggregate, SearchResponse } from '@elastic/elasticsearch/lib/api/types';
+import type {
+  AggregationsAggregate,
+  SearchResponse,
+  SortResults,
+} from '@elastic/elasticsearch/lib/api/types';
 import type { Tier, UsageRecord } from '../types';
 import type { CloudSecurityMeteringCallbackInput } from './types';
 import { CLOUD_DEFEND, CLOUD_SECURITY_TASK_TYPE, CLOUD_DEFEND_HEARTBEAT_INDEX } from './constants';
 
-const BATCH_SIZE = 1000;
+const BATCH_SIZE = 1;
+const SAMPLE_WEIGHT_SECONDS = 3600; // 1 Hour
 
 export interface CloudDefendHeartbeat {
   '@timestamp': string;
@@ -34,13 +39,13 @@ const buildMeteringRecord = (
   timestamp.setMilliseconds(0);
   const creationTimestamp = new Date();
   const usageRecord = {
-    id: `defend-for-containers-${agentId}-${timestamp.toISOString()}`,
+    id: `${projectId}_${agentId}_${timestamp.toISOString()}`,
     usage_timestamp: timestampStr,
     creation_timestamp: creationTimestamp.toISOString(),
     usage: {
       type: CLOUD_SECURITY_TASK_TYPE,
       sub_type: CLOUD_DEFEND,
-      period_seconds: 3600,
+      period_seconds: SAMPLE_WEIGHT_SECONDS,
       quantity: 1,
     },
     source: {
@@ -57,7 +62,7 @@ const buildMeteringRecord = (
 export const getUsageRecords = async (
   esClient: ElasticsearchClient,
   searchFrom: Date,
-  searchAfter?: number[]
+  searchAfter?: SortResults
 ): Promise<SearchResponse<CloudDefendHeartbeat, Record<string, AggregationsAggregate>>> => {
   return esClient.search<CloudDefendHeartbeat>(
     {
@@ -99,7 +104,7 @@ export const getCloudDefendUsageRecords = async ({
 }: CloudSecurityMeteringCallbackInput): Promise<UsageRecord[] | undefined> => {
   try {
     let allRecords: UsageRecord[] = [];
-    let searchAfter: number[] | undefined;
+    let searchAfter: SortResults | undefined;
     let fetchMore = true;
 
     while (fetchMore) {

x-pack/plugins/security_solution_serverless/server/task_manager/usage_reporting_task.ts

@@ -153,8 +153,6 @@ export class SecurityUsageReportingTask {
       return { state: taskInstance.state, runAt: new Date() };
     }
 
-    console.log({ usageRecords });
-
     this.logger.debug(`received usage records: ${JSON.stringify(usageRecords)}`);
 
     let usageReportResponse: Response | undefined;