From 6068fb6472b2fa6b3abba50206600aff802b641a Mon Sep 17 00:00:00 2001
From: Subhobrata Dey <sbcd90@gmail.com>
Date: Thu, 14 Sep 2023 23:55:49 +0000
Subject: [PATCH] fix flaky integration tests

Signed-off-by: Subhobrata Dey <sbcd90@gmail.com>
---
 .../SecurityAnalyticsRestTestCase.java        | 34 +++++--------------
 .../CorrelationEngineRestApiIT.java           |  6 +++-
 2 files changed, 14 insertions(+), 26 deletions(-)

diff --git a/src/test/java/org/opensearch/securityanalytics/SecurityAnalyticsRestTestCase.java b/src/test/java/org/opensearch/securityanalytics/SecurityAnalyticsRestTestCase.java
index 4757a23b5..d3d2192cc 100644
--- a/src/test/java/org/opensearch/securityanalytics/SecurityAnalyticsRestTestCase.java
+++ b/src/test/java/org/opensearch/securityanalytics/SecurityAnalyticsRestTestCase.java
@@ -224,7 +224,7 @@ protected List<Map<String, Object>> searchCorrelatedFindings(String findingId, S
     }
 
     @Before
-    void setDebugLogLevel() throws IOException {
+    void setDebugLogLevel() throws IOException, InterruptedException {
         StringEntity se = new StringEntity("{\n" +
                 "                    \"transient\": {\n" +
                 "                        \"logger.org.opensearch.securityanalytics\":\"DEBUG\",\n" +
@@ -236,6 +236,7 @@ void setDebugLogLevel() throws IOException {
 
 
         makeRequest(client(), "PUT", "_cluster/settings", Collections.emptyMap(), se, new BasicHeader("Content-Type", "application/json"));
+        Thread.sleep(10000);
     }
 
     protected final List<String> clusterPermissions = List.of(
@@ -480,30 +481,13 @@ protected Response refreshIndex(String index) throws IOException {
 
     @SuppressWarnings("unchecked")
     protected List<String> getRandomPrePackagedRules() throws IOException {
-        String request = "{\n" +
-                "  \"from\": 0\n," +
-                "  \"size\": 2000\n," +
-                "  \"query\": {\n" +
-                "    \"nested\": {\n" +
-                "      \"path\": \"rule\",\n" +
-                "      \"query\": {\n" +
-                "        \"bool\": {\n" +
-                "          \"must\": [\n" +
-                "            { \"match\": {\"rule.category\": \"" + TestHelpers.randomDetectorType().toLowerCase(Locale.ROOT) + "\"}}\n" +
-                "          ]\n" +
-                "        }\n" +
-                "      }\n" +
-                "    }\n" +
-                "  }\n" +
-                "}";
-
-        Response searchResponse = makeRequest(client(), "POST", String.format(Locale.getDefault(), "%s/_search", SecurityAnalyticsPlugin.RULE_BASE_URI), Collections.singletonMap("pre_packaged", "true"),
-                new StringEntity(request), new BasicHeader("Content-Type", "application/json"));
-        Assert.assertEquals("Searching rules failed", RestStatus.OK, restStatus(searchResponse));
-
-        Map<String, Object> responseBody = asMap(searchResponse);
-        List<Map<String, Object>> hits = ((List<Map<String, Object>>) ((Map<String, Object>) responseBody.get("hits")).get("hits"));
-        return hits.stream().map(hit -> hit.get("_id").toString()).collect(Collectors.toList());
+        return List.of(
+                "36a037c4-c228-4866-b6a3-48eb292b9955",
+                "c6e91a02-d771-4a6d-a700-42587e0b1095",
+                "5a919691-7302-437f-8e10-1fe088afa145",
+                "e5a6b256-3e47-40fc-89d2-7a477edd6915",
+                "06724b9a-52fc-11ed-bdc3-0242ac120002"
+        );
     }
 
     protected List<String> createAggregationRules () throws IOException {
diff --git a/src/test/java/org/opensearch/securityanalytics/correlation/CorrelationEngineRestApiIT.java b/src/test/java/org/opensearch/securityanalytics/correlation/CorrelationEngineRestApiIT.java
index 4b9a9cb3f..e721e1124 100644
--- a/src/test/java/org/opensearch/securityanalytics/correlation/CorrelationEngineRestApiIT.java
+++ b/src/test/java/org/opensearch/securityanalytics/correlation/CorrelationEngineRestApiIT.java
@@ -92,7 +92,7 @@ public void testBasicCorrelationEngineWorkflow() throws IOException {
     }
 
     @SuppressWarnings("unchecked")
-    public void testListCorrelationsWorkflow() throws IOException {
+    public void testListCorrelationsWorkflow() throws IOException, InterruptedException {
         Long startTime = System.currentTimeMillis();
         LogIndices indices = createIndices();
 
@@ -100,6 +100,7 @@ public void testListCorrelationsWorkflow() throws IOException {
         String testWindowsMonitorId = createTestWindowsDetector(indices.windowsIndex);
 
         createNetworkToAdLdapToWindowsRule(indices);
+        Thread.sleep(5000);
 
         indexDoc(indices.windowsIndex, "2", randomDoc());
         Response executeResponse = executeAlertingMonitor(testWindowsMonitorId, Collections.emptyMap());
@@ -107,11 +108,14 @@ public void testListCorrelationsWorkflow() throws IOException {
         int noOfSigmaRuleMatches = ((List<Map<String, Object>>) ((Map<String, Object>) executeResults.get("input_results")).get("results")).get(0).size();
         Assert.assertEquals(5, noOfSigmaRuleMatches);
 
+        Thread.sleep(5000);
         indexDoc(indices.vpcFlowsIndex, "1", randomVpcFlowDoc());
         executeResponse = executeAlertingMonitor(vpcFlowMonitorId, Collections.emptyMap());
         executeResults = entityAsMap(executeResponse);
         noOfSigmaRuleMatches = ((List<Map<String, Object>>) ((Map<String, Object>) executeResults.get("input_results")).get("results")).get(0).size();
         Assert.assertEquals(1, noOfSigmaRuleMatches);
+
+        Thread.sleep(5000);
         Long endTime = System.currentTimeMillis();
 
         Request request = new Request("GET", "/_plugins/_security_analytics/correlations?start_timestamp=" + startTime + "&end_timestamp=" + endTime);