From 3d21c7b9bf466c901575a37794d9d7bdcc9be52d Mon Sep 17 00:00:00 2001
From: Ido Cohen <90558359+CohenIdo@users.noreply.github.com>
Date: Tue, 25 Feb 2025 13:24:50 +0200
Subject: [PATCH] Fix Chart Trend by Severity Bug

(cherry picked from commit 7f2df7cd427506486ac8f7a34498d0448c431992)
---
 .../server/tasks/findings_stats_task.ts       |  5 +-
 .../server/tasks/task_state.test.ts           | 82 +++++++++++++++++++
 2 files changed, 84 insertions(+), 3 deletions(-)

diff --git a/x-pack/solutions/security/plugins/cloud_security_posture/server/tasks/findings_stats_task.ts b/x-pack/solutions/security/plugins/cloud_security_posture/server/tasks/findings_stats_task.ts
index 673ec5cb9f7a9..1e57b12ea08b1 100644
--- a/x-pack/solutions/security/plugins/cloud_security_posture/server/tasks/findings_stats_task.ts
+++ b/x-pack/solutions/security/plugins/cloud_security_posture/server/tasks/findings_stats_task.ts
@@ -339,7 +339,7 @@ const getFindingsScoresDocIndexingPromises = (
     });
   });
 
-const getVulnStatsTrendDocIndexingPromises = (
+export const getVulnStatsTrendDocIndexingPromises = (
   esClient: ElasticsearchClient,
   vulnStatsAggs?: VulnSeverityAggs
 ) => {
@@ -348,12 +348,11 @@ const getVulnStatsTrendDocIndexingPromises = (
   const scoreByCloudAccount = Object.fromEntries(
     vulnStatsAggs.vulnerabilities_stats_by_cloud_account.buckets.map((accountScore) => {
       const cloudAccountId = accountScore.key;
-
       return [
         cloudAccountId,
         {
           cloudAccountId: accountScore.key,
-          cloudAccountName: accountScore.cloud_account_name.buckets[0].key,
+          cloudAccountName: accountScore.cloud_account_name.buckets[0]?.key || '',
           critical: accountScore.critical.doc_count,
           high: accountScore.high.doc_count,
           medium: accountScore.medium.doc_count,
diff --git a/x-pack/solutions/security/plugins/cloud_security_posture/server/tasks/task_state.test.ts b/x-pack/solutions/security/plugins/cloud_security_posture/server/tasks/task_state.test.ts
index c8289e57b4fbc..c88acb05e40af 100644
--- a/x-pack/solutions/security/plugins/cloud_security_posture/server/tasks/task_state.test.ts
+++ b/x-pack/solutions/security/plugins/cloud_security_posture/server/tasks/task_state.test.ts
@@ -7,6 +7,10 @@
 
 import { cloneDeep } from 'lodash';
 import { stateSchemaByVersion } from './task_state';
+import { getVulnStatsTrendDocIndexingPromises } from './findings_stats_task';
+import { elasticsearchClientMock } from '@kbn/core-elasticsearch-client-server-mocks';
+
+const esClient = elasticsearchClientMock.createClusterClient().asScoped().asInternalUser;
 
 describe('finding stats task state', () => {
   describe('v1', () => {
@@ -35,5 +39,83 @@ describe('finding stats task state', () => {
       const result = v1.up(state);
       expect(result).not.toHaveProperty('foo');
     });
+
+    it('should return undefined if vulnStatsAggs is not provided', () => {
+      const result = getVulnStatsTrendDocIndexingPromises(esClient);
+      expect(result).toBeUndefined();
+      expect(esClient.index).not.toHaveBeenCalled();
+    });
+
+    it('should correctly process vulnerability stats and call esClient.index', async () => {
+      const vulnStatsAggs = {
+        vulnerabilities_stats_by_cloud_account: {
+          buckets: [
+            {
+              key: 'account1',
+              cloud_account_name: { buckets: [{ key: 'Account One' }] },
+              critical: { doc_count: 5 },
+              high: { doc_count: 10 },
+              medium: { doc_count: 15 },
+              low: { doc_count: 20 },
+            },
+          ],
+        },
+        critical: { doc_count: 50 },
+        high: { doc_count: 100 },
+        medium: { doc_count: 150 },
+        low: { doc_count: 200 },
+      };
+
+      await getVulnStatsTrendDocIndexingPromises(esClient, vulnStatsAggs);
+
+      expect(esClient.index).toHaveBeenCalledWith({
+        index: expect.any(String),
+        document: expect.objectContaining({
+          vulnerabilities_stats_by_cloud_account: {
+            account1: expect.objectContaining({
+              cloudAccountName: 'Account One',
+              critical: 5,
+              high: 10,
+              medium: 15,
+              low: 20,
+            }),
+          },
+        }),
+      });
+    });
+
+    it('should handle missing cloud account name gracefully', async () => {
+      const vulnStatsAggs = {
+        vulnerabilities_stats_by_cloud_account: {
+          buckets: [
+            {
+              key: 'account2',
+              cloud_account_name: { buckets: [] },
+              critical: { doc_count: 3 },
+              high: { doc_count: 6 },
+              medium: { doc_count: 9 },
+              low: { doc_count: 12 },
+            },
+          ],
+        },
+        critical: { doc_count: 30 },
+        high: { doc_count: 60 },
+        medium: { doc_count: 90 },
+        low: { doc_count: 120 },
+      };
+
+      await getVulnStatsTrendDocIndexingPromises(esClient, vulnStatsAggs);
+
+      expect(esClient.index).toHaveBeenCalledWith({
+        index: expect.any(String),
+        document: expect.objectContaining({
+          vulnerabilities_stats_by_cloud_account: {
+            account2: expect.objectContaining({
+              cloudAccountName: '',
+            }),
+          },
+        }),
+      });
+    });
   });
 });