From d5aa55159070bd716892360f0321505fb44943c0 Mon Sep 17 00:00:00 2001
From: Ryland Herrick <ryalnd@gmail.com>
Date: Fri, 12 Jul 2024 16:50:42 -0500
Subject: [PATCH] Add tests validating enrichment behavior for ML suppression

This was requested during review of #181926, and I'm circling back to
that now.
---
 .../machine_learning_alert_suppression.ts     | 62 ++++++++++++++++++-
 1 file changed, 61 insertions(+), 1 deletion(-)

diff --git a/x-pack/test/security_solution_api_integration/test_suites/detections_response/detection_engine/rule_execution_logic/trial_license_complete_tier/execution_logic/machine_learning_alert_suppression.ts b/x-pack/test/security_solution_api_integration/test_suites/detections_response/detection_engine/rule_execution_logic/trial_license_complete_tier/execution_logic/machine_learning_alert_suppression.ts
index d8869681de692..fc24c713c99d1 100644
--- a/x-pack/test/security_solution_api_integration/test_suites/detections_response/detection_engine/rule_execution_logic/trial_license_complete_tier/execution_logic/machine_learning_alert_suppression.ts
+++ b/x-pack/test/security_solution_api_integration/test_suites/detections_response/detection_engine/rule_execution_logic/trial_license_complete_tier/execution_logic/machine_learning_alert_suppression.ts
@@ -22,7 +22,10 @@ import {
   TIMESTAMP,
 } from '@kbn/rule-data-utils';
 import { ALERT_ORIGINAL_TIME } from '@kbn/security-solution-plugin/common/field_maps/field_names';
-import { DETECTION_ENGINE_SIGNALS_STATUS_URL as DETECTION_ENGINE_ALERTS_STATUS_URL } from '@kbn/security-solution-plugin/common/constants';
+import {
+  DETECTION_ENGINE_SIGNALS_STATUS_URL as DETECTION_ENGINE_ALERTS_STATUS_URL,
+  ENABLE_ASSET_CRITICALITY_SETTING,
+} from '@kbn/security-solution-plugin/common/constants';
 import { EsArchivePathBuilder } from '../../../../../../es_archive_path_builder';
 import { FtrProviderContext } from '../../../../../../ftr_provider_context';
 import {
@@ -1102,6 +1105,63 @@ export default ({ getService }: FtrProviderContext) => {
           });
         });
       });
+
+      describe('with enrichments', () => {
+        const kibanaServer = getService('kibanaServer');
+
+        before(async () => {
+          await esArchiver.load('x-pack/test/functional/es_archives/entity/risks');
+          await esArchiver.load('x-pack/test/functional/es_archives/asset_criticality');
+          await kibanaServer.uiSettings.update({
+            [ENABLE_ASSET_CRITICALITY_SETTING]: true,
+          });
+        });
+
+        after(async () => {
+          await esArchiver.unload('x-pack/test/functional/es_archives/entity/risks');
+          await esArchiver.unload('x-pack/test/functional/es_archives/asset_criticality');
+        });
+
+        beforeEach(async () => {
+          const timestamp = new Date().toISOString();
+          const anomalyWithKnownEntities = {
+            ...baseAnomaly,
+            timestamp,
+            user: { name: 'root' },
+            host: { name: 'zeek-newyork-sha-aa8df15' },
+          };
+          await indexListOfDocuments([anomalyWithKnownEntities]);
+
+          ruleProps = {
+            ...baseRuleProps,
+            from: timestamp,
+            alert_suppression: {
+              group_by: ['host.name'],
+              missing_fields_strategy: 'suppress',
+            },
+          };
+        });
+
+        it('should be enriched with host risk score', async () => {
+          const { previewId } = await previewRule({ supertest, rule: ruleProps });
+          const previewAlerts = await getPreviewAlerts({ es, previewId });
+          expect(previewAlerts).toHaveLength(1);
+          const alertSource = previewAlerts[0]._source;
+
+          expect(alertSource?.host?.risk?.calculated_level).toBe('Low');
+          expect(alertSource?.host?.risk?.calculated_score_norm).toBe(23);
+        });
+
+        it('should be enriched alert with criticality_level', async () => {
+          const { previewId } = await previewRule({ supertest, rule: ruleProps });
+          const previewAlerts = await getPreviewAlerts({ es, previewId });
+          expect(previewAlerts).toHaveLength(1);
+          const fullAlert = previewAlerts[0]._source;
+
+          expect(fullAlert?.['host.asset.criticality']).toBe('medium_impact');
+          expect(fullAlert?.['user.asset.criticality']).toBe('extreme_impact');
+        });
+      });
     });
   });
 };